Compare commits

...

2 Commits

Author SHA1 Message Date
HAL9000 86a990e533 fix(context): fix Behave table format and lint issues in dynamic budget allocation tests
CI / lint (pull_request) Failing after 0s
CI / typecheck (pull_request) Failing after 0s
CI / security (pull_request) Failing after 0s
CI / quality (pull_request) Failing after 0s
CI / coverage (pull_request) Has been skipped
CI / unit_tests (pull_request) Failing after 0s
CI / docker (pull_request) Has been skipped
CI / integration_tests (pull_request) Failing after 0s
CI / push-validation (pull_request) Failing after 0s
CI / helm (pull_request) Successful in 31s
CI / build (pull_request) Successful in 37s
CI / e2e_tests (pull_request) Successful in 3m13s
CI / status-check (pull_request) Failing after 0s
2026-04-24 07:13:33 +00:00
HAL9000 54589a830a feat(context): implement dynamic budget allocation for task-complexity-aware context assembly
CI / push-validation (pull_request) Successful in 25s
CI / lint (pull_request) Failing after 51s
CI / helm (pull_request) Successful in 32s
CI / build (pull_request) Successful in 3m49s
CI / quality (pull_request) Successful in 4m28s
CI / typecheck (pull_request) Successful in 4m36s
CI / security (pull_request) Successful in 5m5s
CI / coverage (pull_request) Has been skipped
CI / unit_tests (pull_request) Failing after 5m37s
CI / docker (pull_request) Has been skipped
CI / e2e_tests (pull_request) Successful in 7m5s
CI / integration_tests (pull_request) Successful in 8m0s
CI / status-check (pull_request) Failing after 4s
2026-04-19 01:12:57 +00:00
3 changed files with 687 additions and 0 deletions
@@ -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
@@ -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,
}
@@ -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.
"""
...