diff --git a/features/context_dynamic_budget_allocation.feature b/features/context_dynamic_budget_allocation.feature new file mode 100644 index 000000000..d55fcfd71 --- /dev/null +++ b/features/context_dynamic_budget_allocation.feature @@ -0,0 +1,139 @@ +Feature: Dynamic budget allocation for task-complexity-aware context assembly + As a context assembly system + I want to dynamically allocate context budget based on task complexity + So that simple tasks don't over-consume context while complex tasks get needed resources + + Background: + Given a DynamicBudgetAllocator with default configuration + And the default configuration has min_budget_tokens of 512 + And the default configuration has max_budget_tokens of 8192 + + Scenario: Simple task with minimal complexity signals + Given complexity signals with: + | key | value | + | prompt_token_count | 100 | + | referenced_file_count | 1 | + | plan_step_count | 1 | + When I allocate budget for these signals + Then the allocated_budget should be between 512 and 1024 + And the complexity_score should be less than 0.3 + + Scenario: Complex task with high complexity signals + Given complexity signals with: + | key | value | + | prompt_token_count | 1500 | + | referenced_file_count | 30 | + | plan_step_count | 15 | + When I allocate budget for these signals + Then the allocated_budget should be between 6000 and 8192 + And the complexity_score should be greater than 0.6 + + Scenario: Medium complexity task + Given complexity signals with: + | key | value | + | prompt_token_count | 500 | + | referenced_file_count | 10 | + | plan_step_count | 5 | + When I allocate budget for these signals + Then the allocated_budget should be between 2000 and 5000 + And the complexity_score should be between 0.3 and 0.6 + + Scenario: Budget allocation respects minimum bound + Given complexity signals with: + | key | value | + | prompt_token_count | 0 | + | referenced_file_count | 0 | + | plan_step_count | 0 | + When I allocate budget for these signals + Then the allocated_budget should be exactly 512 + + Scenario: Budget allocation respects maximum bound + Given complexity signals with: + | key | value | + | prompt_token_count | 5000 | + | referenced_file_count | 100 | + | plan_step_count | 50 | + When I allocate budget for these signals + Then the allocated_budget should be exactly 8192 + + Scenario: Historical average influences budget allocation + Given complexity signals with: + | key | value | + | prompt_token_count | 300 | + | referenced_file_count | 5 | + | plan_step_count | 3 | + | historical_average_tokens | 3000 | + When I allocate budget for these signals + Then the complexity_score should be greater than 0.3 + And the signal_contributions should include historical contribution + + Scenario: Configuration with custom weights + Given a DynamicBudgetAllocator with custom configuration: + | key | value | + | min_budget_tokens | 256 | + | max_budget_tokens | 4096 | + | prompt_weight | 0.5 | + | file_count_weight | 0.2 | + | plan_depth_weight | 0.2 | + | historical_weight | 0.1 | + And complexity signals with: + | key | value | + | prompt_token_count | 1000 | + | referenced_file_count | 10 | + | plan_step_count | 5 | + When I allocate budget for these signals + Then the allocated_budget should be between 256 and 4096 + And the signal_contributions["prompt"] should be greater than signal_contributions["file_count"] + + Scenario: Invalid complexity signals are rejected + Given complexity signals with: + | key | value | + | prompt_token_count | -1 | + | referenced_file_count | 5 | + | plan_step_count | 3 | + When I try to allocate budget for these signals + Then an error should be raised with message containing "non-negative" + + Scenario: Invalid configuration is rejected + When I try to create a DynamicBudgetAllocator with invalid configuration: + | key | value | + | min_budget_tokens | 1000 | + | max_budget_tokens | 500 | + Then an error should be raised with message containing "max_budget_tokens" + + Scenario: Reasoning explains allocation decision + Given complexity signals with: + | key | value | + | prompt_token_count | 500 | + | referenced_file_count | 10 | + | plan_step_count | 5 | + When I allocate budget for these signals + Then the reasoning should contain "Dynamic Budget Allocation Analysis" + And the reasoning should contain "complexity score" + And the reasoning should contain "Allocated budget" + + Scenario: Integration with context assembler + Given a context assembler with dynamic budget allocation + And a plan with complexity signals: + | key | value | + | prompt_token_count | 400 | + | referenced_file_count | 8 | + | plan_step_count | 4 | + When the context assembler assembles context + Then the context should use the dynamically allocated budget + And the budget should be between 512 and 8192 + + Scenario: Simple vs complex task budget comparison + Given a simple task with complexity signals: + | key | value | + | prompt_token_count | 100 | + | referenced_file_count | 1 | + | plan_step_count | 1 | + And a complex task with complexity signals: + | key | value | + | prompt_token_count | 1500 | + | referenced_file_count | 30 | + | plan_step_count | 15 | + When I allocate budget for both tasks + Then the complex task budget should be significantly higher than simple task budget + And the complex task budget should be at least 2x the simple task budget diff --git a/features/steps/context_dynamic_budget_allocation_steps.py b/features/steps/context_dynamic_budget_allocation_steps.py new file mode 100644 index 000000000..baedc9ee9 --- /dev/null +++ b/features/steps/context_dynamic_budget_allocation_steps.py @@ -0,0 +1,272 @@ +"""Step definitions for dynamic budget allocation feature.""" + +from __future__ import annotations + +from behave import given, then, when +from behave.runner import Context + +from cleveragents.application.services.dynamic_budget_allocator import ( + BudgetAllocationConfig, + ComplexitySignals, + DynamicBudgetAllocator, +) + + +@given("a DynamicBudgetAllocator with default configuration") +def step_create_default_allocator(context: Context) -> None: + """Create allocator with default configuration.""" + context.allocator = DynamicBudgetAllocator() + context.config = context.allocator.config + + +@given("the default configuration has min_budget_tokens of {value:d}") +def step_check_min_budget(context: Context, value: int) -> None: + """Verify minimum budget configuration.""" + assert context.config.min_budget_tokens == value + + +@given("the default configuration has max_budget_tokens of {value:d}") +def step_check_max_budget(context: Context, value: int) -> None: + """Verify maximum budget configuration.""" + assert context.config.max_budget_tokens == value + + +@given("complexity signals with:") +def step_create_complexity_signals(context: Context) -> None: + """Create complexity signals from table.""" + signals_dict = {} + for row in context.table: + signals_dict[row["key"]] = int(row["value"]) + + context.signals = ComplexitySignals( + prompt_token_count=signals_dict.get("prompt_token_count", 0), + referenced_file_count=signals_dict.get("referenced_file_count", 0), + plan_step_count=signals_dict.get("plan_step_count", 0), + historical_average_tokens=signals_dict.get("historical_average_tokens"), + ) + + +@when("I allocate budget for these signals") +def step_allocate_budget(context: Context) -> None: + """Allocate budget using the allocator.""" + context.result = context.allocator.allocate(context.signals) + + +@then("the allocated_budget should be between {min_val:d} and {max_val:d}") +def step_check_budget_range(context: Context, min_val: int, max_val: int) -> None: + """Verify allocated budget is within range.""" + assert min_val <= context.result.allocated_budget <= max_val, ( + f"Budget {context.result.allocated_budget} not in range [{min_val}, {max_val}]" + ) + + +@then("the allocated_budget should be exactly {value:d}") +def step_check_budget_exact(context: Context, value: int) -> None: + """Verify allocated budget is exact value.""" + assert context.result.allocated_budget == value + + +@then("the complexity_score should be less than {value:f}") +def step_check_score_less(context: Context, value: float) -> None: + """Verify complexity score is less than value.""" + assert context.result.complexity_score < value + + +@then("the complexity_score should be greater than {value:f}") +def step_check_score_greater(context: Context, value: float) -> None: + """Verify complexity score is greater than value.""" + assert context.result.complexity_score > value + + +@then("the complexity_score should be between {min_val:f} and {max_val:f}") +def step_check_score_range(context: Context, min_val: float, max_val: float) -> None: + """Verify complexity score is within range.""" + assert min_val <= context.result.complexity_score <= max_val, ( + f"Score {context.result.complexity_score} not in range [{min_val}, {max_val}]" + ) + + +@then("the signal_contributions should include historical contribution") +def step_check_historical_contribution(context: Context) -> None: + """Verify historical signal contribution exists.""" + assert "historical" in context.result.signal_contributions + assert context.result.signal_contributions["historical"] > 0 + + +@given("a DynamicBudgetAllocator with custom configuration:") +def step_create_custom_allocator(context: Context) -> None: + """Create allocator with custom configuration.""" + config_dict = {} + for row in context.table: + k = row["key"] + v = row["value"] + if k in ["min_budget_tokens", "max_budget_tokens"]: + config_dict[k] = int(v) + else: + config_dict[k] = float(v) + + config = BudgetAllocationConfig(**config_dict) + context.allocator = DynamicBudgetAllocator(config) + + +@then( + "the signal_contributions[{key}] should be greater than signal_contributions[{other_key}]" +) +def step_compare_contributions(context: Context, key: str, other_key: str) -> None: + """Compare signal contributions.""" + key = key.strip("'\"") + other_key = other_key.strip("'\"") + assert ( + context.result.signal_contributions[key] + > context.result.signal_contributions[other_key] + ) + + +@when("I try to allocate budget for these signals") +def step_try_allocate_budget(context: Context) -> None: + """Try to allocate budget and catch errors.""" + try: + context.result = context.allocator.allocate(context.signals) + context.error = None + except Exception as e: + context.error = e + + +@then("an error should be raised with message containing {text}") +def step_check_error_message(context: Context, text: str) -> None: + """Verify error message contains text.""" + text = text.strip(chr(34)) + assert context.error is not None + assert text in str(context.error) + + +@when("I try to create a DynamicBudgetAllocator with invalid configuration:") +def step_try_create_invalid_allocator(context: Context) -> None: + """Try to create allocator with invalid config.""" + config_dict = {} + for row in context.table: + config_dict[row["key"]] = int(row["value"]) + + try: + BudgetAllocationConfig(**config_dict) + context.error = None + except Exception as e: + context.error = e + + +@then("the reasoning should contain {text}") +def step_check_reasoning(context: Context, text: str) -> None: + """Verify reasoning contains text.""" + text = text.strip(chr(34)) + assert text in context.result.reasoning + + +@given("a context assembler with dynamic budget allocation") +def step_create_context_assembler(context: Context) -> None: + """Create a mock context assembler.""" + context.assembler = MockContextAssembler(context.allocator) + + +@given("a plan with complexity signals:") +def step_create_plan_with_signals(context: Context) -> None: + """Create a plan with complexity signals.""" + signals_dict = {} + for row in context.table: + signals_dict[row["key"]] = int(row["value"]) + + context.plan_signals = ComplexitySignals( + prompt_token_count=signals_dict.get("prompt_token_count", 0), + referenced_file_count=signals_dict.get("referenced_file_count", 0), + plan_step_count=signals_dict.get("plan_step_count", 0), + ) + + +@when("the context assembler assembles context") +def step_assemble_context(context: Context) -> None: + """Assemble context using the assembler.""" + context.assembled_context = context.assembler.assemble(context.plan_signals) + + +@then("the context should use the dynamically allocated budget") +def step_check_context_uses_budget(context: Context) -> None: + """Verify context uses allocated budget.""" + assert context.assembled_context["budget"] > 0 + + +@then("the budget should be between {min_val:d} and {max_val:d}") +def step_check_context_budget_range( + context: Context, min_val: int, max_val: int +) -> None: + """Verify context budget is within range.""" + budget = context.assembled_context["budget"] + assert min_val <= budget <= max_val + + +@given("a simple task with complexity signals:") +def step_create_simple_task(context: Context) -> None: + """Create a simple task.""" + signals_dict = {} + for row in context.table: + signals_dict[row["key"]] = int(row["value"]) + + context.simple_signals = ComplexitySignals( + prompt_token_count=signals_dict.get("prompt_token_count", 0), + referenced_file_count=signals_dict.get("referenced_file_count", 0), + plan_step_count=signals_dict.get("plan_step_count", 0), + ) + + +@given("a complex task with complexity signals:") +def step_create_complex_task(context: Context) -> None: + """Create a complex task.""" + signals_dict = {} + for row in context.table: + signals_dict[row["key"]] = int(row["value"]) + + context.complex_signals = ComplexitySignals( + prompt_token_count=signals_dict.get("prompt_token_count", 0), + referenced_file_count=signals_dict.get("referenced_file_count", 0), + plan_step_count=signals_dict.get("plan_step_count", 0), + ) + + +@when("I allocate budget for both tasks") +def step_allocate_both_tasks(context: Context) -> None: + """Allocate budget for both simple and complex tasks.""" + context.simple_result = context.allocator.allocate(context.simple_signals) + context.complex_result = context.allocator.allocate(context.complex_signals) + + +@then("the complex task budget should be significantly higher than simple task budget") +def step_check_budget_difference(context: Context) -> None: + """Verify complex task has higher budget.""" + assert ( + context.complex_result.allocated_budget > context.simple_result.allocated_budget + ) + + +@then( + "the complex task budget should be at least {multiplier:d}x the simple task budget" +) +def step_check_budget_multiplier(context: Context, multiplier: int) -> None: + """Verify complex task budget is at least N times simple task budget.""" + assert ( + context.complex_result.allocated_budget + >= context.simple_result.allocated_budget * multiplier + ) + + +class MockContextAssembler: + """Mock context assembler for testing.""" + + def __init__(self, allocator: DynamicBudgetAllocator) -> None: + """Initialize with allocator.""" + self.allocator = allocator + + def assemble(self, signals: ComplexitySignals) -> dict: + """Assemble context with dynamic budget.""" + result = self.allocator.allocate(signals) + return { + "budget": result.allocated_budget, + "complexity_score": result.complexity_score, + } diff --git a/src/cleveragents/application/services/dynamic_budget_allocator.py b/src/cleveragents/application/services/dynamic_budget_allocator.py new file mode 100644 index 000000000..f500605b8 --- /dev/null +++ b/src/cleveragents/application/services/dynamic_budget_allocator.py @@ -0,0 +1,276 @@ +"""Dynamic budget allocation algorithm for task-complexity-aware context assembly. + +This module implements a DynamicBudgetAllocator that analyzes task complexity signals +and adjusts the context budget accordingly. It ensures efficient use of the LLM's +context window across diverse task types. + +The allocator considers: +- Prompt token count: Size of the initial prompt +- Referenced file count: Number of files in context +- Plan step count: Depth/complexity of the execution plan +- Historical average: Average token usage from past executions +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Protocol + +import structlog + +logger = structlog.get_logger(__name__) + + +@dataclass +class ComplexitySignals: + """Signals used to determine task complexity.""" + + prompt_token_count: int + """Number of tokens in the initial prompt.""" + + referenced_file_count: int + """Number of files referenced in the context.""" + + plan_step_count: int + """Number of steps in the execution plan.""" + + historical_average_tokens: int | None = None + """Average token usage from historical executions.""" + + def __post_init__(self) -> None: + """Validate signal values.""" + if self.prompt_token_count < 0: + raise ValueError("prompt_token_count must be non-negative") + if self.referenced_file_count < 0: + raise ValueError("referenced_file_count must be non-negative") + if self.plan_step_count < 0: + raise ValueError("plan_step_count must be non-negative") + if ( + self.historical_average_tokens is not None + and self.historical_average_tokens < 0 + ): + raise ValueError("historical_average_tokens must be non-negative") + + +@dataclass +class BudgetAllocationConfig: + """Configuration for dynamic budget allocation.""" + + min_budget_tokens: int = 512 + """Minimum context budget in tokens.""" + + max_budget_tokens: int = 8192 + """Maximum context budget in tokens.""" + + prompt_weight: float = 0.3 + """Weight for prompt token count signal (0.0 to 1.0).""" + + file_count_weight: float = 0.25 + """Weight for referenced file count signal (0.0 to 1.0).""" + + plan_depth_weight: float = 0.25 + """Weight for plan step count signal (0.0 to 1.0).""" + + historical_weight: float = 0.2 + """Weight for historical average signal (0.0 to 1.0).""" + + file_count_scaling_factor: float = 100.0 + """Tokens per file for scaling file count signal.""" + + plan_depth_scaling_factor: float = 50.0 + """Tokens per plan step for scaling plan depth signal.""" + + def __post_init__(self) -> None: + """Validate configuration values.""" + if self.min_budget_tokens < 0: + raise ValueError("min_budget_tokens must be non-negative") + if self.max_budget_tokens < self.min_budget_tokens: + raise ValueError("max_budget_tokens must be >= min_budget_tokens") + if not (0.0 <= self.prompt_weight <= 1.0): + raise ValueError("prompt_weight must be between 0.0 and 1.0") + if not (0.0 <= self.file_count_weight <= 1.0): + raise ValueError("file_count_weight must be between 0.0 and 1.0") + if not (0.0 <= self.plan_depth_weight <= 1.0): + raise ValueError("plan_depth_weight must be between 0.0 and 1.0") + if not (0.0 <= self.historical_weight <= 1.0): + raise ValueError("historical_weight must be between 0.0 and 1.0") + if self.file_count_scaling_factor <= 0: + raise ValueError("file_count_scaling_factor must be positive") + if self.plan_depth_scaling_factor <= 0: + raise ValueError("plan_depth_scaling_factor must be positive") + + +@dataclass +class BudgetAllocationResult: + """Result of budget allocation.""" + + allocated_budget: int + """The allocated context budget in tokens.""" + + complexity_score: float + """Normalized complexity score (0.0 to 1.0).""" + + signal_contributions: dict[str, float] = field(default_factory=dict) + """Contribution of each signal to the final budget.""" + + reasoning: str = "" + """Human-readable explanation of the allocation.""" + + +class DynamicBudgetAllocator: + """Allocates context budget based on task complexity signals. + + This allocator analyzes multiple complexity signals and computes a weighted + complexity score. The score is then used to dynamically adjust the context + budget within configured min/max bounds. + + Example: + >>> config = BudgetAllocationConfig( + ... min_budget_tokens=512, + ... max_budget_tokens=8192, + ... ) + >>> allocator = DynamicBudgetAllocator(config) + >>> signals = ComplexitySignals( + ... prompt_token_count=500, + ... referenced_file_count=10, + ... plan_step_count=5, + ... ) + >>> result = allocator.allocate(signals) + >>> print(f"Allocated budget: {result.allocated_budget} tokens") + """ + + def __init__(self, config: BudgetAllocationConfig | None = None) -> None: + """Initialize the allocator with configuration. + + Args: + config: Configuration for budget allocation. If None, uses defaults. + """ + self.config = config or BudgetAllocationConfig() + self._logger = logger.bind(component="dynamic_budget_allocator") + + def allocate(self, signals: ComplexitySignals) -> BudgetAllocationResult: + """Allocate context budget based on complexity signals. + + Args: + signals: Complexity signals for the task. + + Returns: + BudgetAllocationResult containing the allocated budget and metadata. + """ + # Compute normalized signal contributions + signal_contributions: dict[str, float] = {} + + # Normalize prompt token count (0-1 scale, assuming max 2000 tokens) + prompt_score = min(1.0, signals.prompt_token_count / 2000.0) + signal_contributions["prompt"] = prompt_score * self.config.prompt_weight + + # Normalize file count (0-1 scale, assuming max 50 files) + file_score = min(1.0, signals.referenced_file_count / 50.0) + signal_contributions["file_count"] = file_score * self.config.file_count_weight + + # Normalize plan depth (0-1 scale, assuming max 20 steps) + plan_score = min(1.0, signals.plan_step_count / 20.0) + signal_contributions["plan_depth"] = plan_score * self.config.plan_depth_weight + + # Normalize historical average (0-1 scale, assuming max 4096 tokens) + historical_score = 0.0 + if signals.historical_average_tokens is not None: + historical_score = min(1.0, signals.historical_average_tokens / 4096.0) + signal_contributions["historical"] = ( + historical_score * self.config.historical_weight + ) + + # Compute weighted complexity score + complexity_score = sum(signal_contributions.values()) + + # Clamp complexity score to [0, 1] + complexity_score = min(1.0, max(0.0, complexity_score)) + + # Allocate budget based on complexity score + budget_range = self.config.max_budget_tokens - self.config.min_budget_tokens + allocated_budget = int( + self.config.min_budget_tokens + (complexity_score * budget_range) + ) + + # Build reasoning string + reasoning = self._build_reasoning( + signals, signal_contributions, complexity_score + ) + + self._logger.info( + "budget_allocated", + prompt_tokens=signals.prompt_token_count, + file_count=signals.referenced_file_count, + plan_steps=signals.plan_step_count, + complexity_score=complexity_score, + allocated_budget=allocated_budget, + ) + + return BudgetAllocationResult( + allocated_budget=allocated_budget, + complexity_score=complexity_score, + signal_contributions=signal_contributions, + reasoning=reasoning, + ) + + def _build_reasoning( + self, + signals: ComplexitySignals, + contributions: dict[str, float], + complexity_score: float, + ) -> str: + """Build a human-readable explanation of the allocation. + + Args: + signals: The complexity signals. + contributions: Signal contributions to the score. + complexity_score: The final complexity score. + + Returns: + A formatted reasoning string. + """ + lines = [ + "Dynamic Budget Allocation Analysis:", + f" Prompt tokens: {signals.prompt_token_count} " + f"(contribution: {contributions['prompt']:.3f})", + f" Referenced files: {signals.referenced_file_count} " + f"(contribution: {contributions['file_count']:.3f})", + f" Plan steps: {signals.plan_step_count} " + f"(contribution: {contributions['plan_depth']:.3f})", + ] + + if signals.historical_average_tokens is not None: + lines.append( + f" Historical average: {signals.historical_average_tokens} tokens " + f"(contribution: {contributions['historical']:.3f})" + ) + + lines.extend( + [ + f" Overall complexity score: {complexity_score:.3f}", + f" Allocated budget: {self.config.min_budget_tokens} - " + f"{self.config.max_budget_tokens} tokens", + ] + ) + + return "\n".join(lines) + + +class ContextAssemblerProtocol(Protocol): + """Protocol for context assemblers that can use dynamic budget allocation.""" + + def assemble( + self, + plan: Any, + budget_allocator: DynamicBudgetAllocator | None = None, + ) -> Any: + """Assemble context with optional dynamic budget allocation. + + Args: + plan: The execution plan. + budget_allocator: Optional allocator for dynamic budgeting. + + Returns: + Assembled context. + """ + ...