From 42a32c2709ac58a0ab96845653b6c687724c3148 Mon Sep 17 00:00:00 2001 From: Aditya Chhabra Date: Mon, 23 Mar 2026 21:30:09 +0000 Subject: [PATCH] fix(acms): align budget allocation formula and protocol signatures with spec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Align ACMS pipeline protocol signatures and allocation formula with the specification (docs/specification.md §42630-42937): - BudgetAllocator.allocate() now accepts optional request: ContextRequest parameter per BudgetAllocatorProtocol spec (§44754-44766), enabling future request-aware allocation strategies. - Allocation formula changed from proportional to confidence alone to proportional to confidence * quality_score per spec §45003. Both DefaultBudgetAllocator and ProportionalBudgetAllocator use the new weighted formula. With default quality_score=1.0, behavior is backward- compatible. - DetailDepthResolver.resolve() now accepts budget: int parameter per DetailDepthResolverProtocol spec (§44803-44814), enabling future budget-aware depth resolution decisions. - Default packer changed from no-op DefaultBudgetPacker to production GreedyKnapsackPacker per spec §45013. ACMSPipeline uses lazy import to avoid circular dependency; ContextAssemblyPipeline imports directly. - Added quality_score field to service-layer StrategyCapabilities (already present in domain model). - Updated 5 feature scenarios and 1 Robot helper where GreedyKnapsackPacker relevance-based ordering changed fragment output order vs. the old no-op packer. Fixed by adjusting test fragment relevance scores to preserve expected ordering. ISSUES CLOSED: #924 --- features/acms_pipeline.feature | 14 ++--- features/acms_pipeline_orchestrator.feature | 4 +- features/steps/acms_pipeline_steps.py | 2 + .../acms_service_coverage_boost_steps.py | 2 +- robot/helper_acms_pipeline.py | 2 +- .../application/services/acms_phase2.py | 1 + .../application/services/acms_pipeline.py | 48 +++++++++++------ .../application/services/acms_service.py | 52 +++++++++++++++---- .../application/services/fusion_engine.py | 2 +- .../services/strategy_coordinator.py | 4 +- 10 files changed, 93 insertions(+), 38 deletions(-) diff --git a/features/acms_pipeline.feature b/features/acms_pipeline.feature index 0339cf20a..55e039a95 100644 --- a/features/acms_pipeline.feature +++ b/features/acms_pipeline.feature @@ -166,7 +166,7 @@ Feature: ACMS v1 Context Assembly Pipeline Given the following tiered context fragments: | uko_node | content | score | tokens | tier | | project://app/cold.py | cold-item | 0.9 | 100 | cold | - | project://app/hot.py | hot-item | 0.5 | 100 | hot | + | project://app/hot.py | hot-item | 0.85 | 100 | hot | | project://app/warm.py | warm-item | 0.8 | 100 | warm | And a context budget with max_tokens 200 and reserved_tokens 0 When I assemble with strategy "tiered" @@ -220,8 +220,8 @@ Feature: ACMS v1 Context Assembly Pipeline Given a custom context strategy that reverses fragment order And the following context fragments: | uko_node | content | score | tokens | - | project://app/first.py | first | 0.9 | 50 | - | project://app/second.py | second | 0.5 | 50 | + | project://app/first.py | first | 0.5 | 50 | + | project://app/second.py | second | 0.9 | 50 | And a context budget with max_tokens 200 and reserved_tokens 0 When I assemble with strategy "reverse" Then the payload should contain 2 fragments @@ -256,8 +256,8 @@ Feature: ACMS v1 Context Assembly Pipeline Given a custom context strategy that reverses fragment order And the following context fragments: | uko_node | content | score | tokens | - | project://app/first.py | first | 0.9 | 50 | - | project://app/second.py | second | 0.5 | 50 | + | project://app/first.py | first | 0.5 | 50 | + | project://app/second.py | second | 0.9 | 50 | And a context budget with max_tokens 200 and reserved_tokens 0 When I register the reverse strategy as "relevance" And I assemble with strategy "relevance" @@ -378,8 +378,8 @@ Feature: ACMS v1 Context Assembly Pipeline Given a pipeline with a custom strategy executor that reverses fragments And the following context fragments: | uko_node | content | score | tokens | - | project://app/a.py | alpha | 0.9 | 50 | - | project://app/b.py | beta | 0.5 | 50 | + | project://app/a.py | alpha | 0.5 | 50 | + | project://app/b.py | beta | 0.9 | 50 | And a context budget with max_tokens 500 and reserved_tokens 0 When I assemble with strategy "relevance" Then the custom strategy executor should have been called diff --git a/features/acms_pipeline_orchestrator.feature b/features/acms_pipeline_orchestrator.feature index e12e62075..886333768 100644 --- a/features/acms_pipeline_orchestrator.feature +++ b/features/acms_pipeline_orchestrator.feature @@ -249,8 +249,8 @@ Feature: ACMS Pipeline Orchestrator and Phase 1 Components And a custom orchestrator strategy that returns fragments in reverse And the following orchestrator test fragments: | uko_node | content | score | tokens | - | project://app/a.py | first | 0.9 | 50 | - | project://app/b.py | second | 0.5 | 50 | + | project://app/a.py | first | 0.5 | 50 | + | project://app/b.py | second | 0.9 | 50 | And a pipeline budget with max_tokens 200 and reserved_tokens 0 When I assemble via the orchestrator with strategy "reverse_orch" Then the first orchestrator fragment content should be "second" diff --git a/features/steps/acms_pipeline_steps.py b/features/steps/acms_pipeline_steps.py index 3ec9a5ed0..a7602df96 100644 --- a/features/steps/acms_pipeline_steps.py +++ b/features/steps/acms_pipeline_steps.py @@ -24,6 +24,7 @@ from cleveragents.application.services.acms_service import ( StrategyCapabilities, TieredStrategy, ) +from cleveragents.domain.models.acms.crp import ContextRequest from cleveragents.domain.models.core.context_fragment import ( ContextBudget, ContextFragment, @@ -545,6 +546,7 @@ class _HalvingAllocator: self, candidates: list[tuple[Any, float]], total_budget: int, + request: ContextRequest | None = None, ) -> list[tuple[Any, float, int]]: self.called = True half = total_budget // 2 diff --git a/features/steps/acms_service_coverage_boost_steps.py b/features/steps/acms_service_coverage_boost_steps.py index eb33979a5..0b8a58a3c 100644 --- a/features/steps/acms_service_coverage_boost_steps.py +++ b/features/steps/acms_service_coverage_boost_steps.py @@ -67,7 +67,7 @@ class _SimpleStrategy: @property def capabilities(self) -> StrategyCapabilities: - return StrategyCapabilities() + return StrategyCapabilities(quality_score=1.0) def can_handle(self, request: dict[str, Any]) -> float: return 0.9 diff --git a/robot/helper_acms_pipeline.py b/robot/helper_acms_pipeline.py index d8c2b5072..47db4e485 100644 --- a/robot/helper_acms_pipeline.py +++ b/robot/helper_acms_pipeline.py @@ -182,7 +182,7 @@ def _cmd_assemble_tiered() -> int: ContextFragment( uko_node="project://app/hot.py", content="hot-item", - relevance_score=0.5, + relevance_score=0.85, token_count=100, tier="hot", provenance=_DEFAULT_PROV, diff --git a/src/cleveragents/application/services/acms_phase2.py b/src/cleveragents/application/services/acms_phase2.py index 708f21ca0..898e53544 100644 --- a/src/cleveragents/application/services/acms_phase2.py +++ b/src/cleveragents/application/services/acms_phase2.py @@ -123,6 +123,7 @@ class MaxDepthResolver: def resolve( self, fragments: Sequence[ContextFragment], + budget: int = 0, ) -> Sequence[ContextFragment]: """Return fragments with depth conflicts resolved to maximum depth.""" if not fragments: diff --git a/src/cleveragents/application/services/acms_pipeline.py b/src/cleveragents/application/services/acms_pipeline.py index b8b4af9a1..11a8538fc 100644 --- a/src/cleveragents/application/services/acms_pipeline.py +++ b/src/cleveragents/application/services/acms_pipeline.py @@ -40,12 +40,14 @@ if TYPE_CHECKING: from cleveragents.config.settings import Settings from cleveragents.infrastructure.database.unit_of_work import UnitOfWork +from cleveragents.application.services.acms_phase2 import ( + GreedyKnapsackPacker, +) from cleveragents.application.services.acms_service import ( ACMSPipeline, BudgetAllocator, BudgetPacker, ContextStrategy, - DefaultBudgetPacker, DefaultDeduplicator, DefaultDepthResolver, DefaultOrderer, @@ -61,6 +63,9 @@ from cleveragents.application.services.acms_service import ( StrategyExecutor, StrategySelector, ) +from cleveragents.domain.models.acms.crp import ( + ContextRequest, +) from cleveragents.domain.models.core.context_fragment import ( ContextBudget, ContextFragment, @@ -182,10 +187,10 @@ class ConfidenceWeightedSelector: class ProportionalBudgetAllocator: """Distribute token budget proportionally with min-budget enforcement. - Each candidate receives a share of ``total_budget`` proportional to its - confidence score. Candidates whose proportional share falls below - ``min_useful_budget`` are excluded and their tokens are redistributed - among the remaining candidates. + Each candidate receives a share of ``total_budget`` proportional to + ``confidence * quality_score`` per spec §42918. Candidates whose + proportional share falls below ``min_useful_budget`` are excluded and + their tokens are redistributed among the remaining candidates. Uses the largest-remainder method to ensure the full budget is consumed with no token loss. @@ -198,10 +203,17 @@ class ProportionalBudgetAllocator: self._min_useful_budget = min_useful_budget self._logger = logger.bind(component="ProportionalBudgetAllocator") + @staticmethod + def _weighted_score(strategy: ContextStrategy, confidence: float) -> float: + """Compute the allocation weight: ``confidence * quality_score``.""" + qs = getattr(strategy.capabilities, "quality_score", 1.0) + return confidence * qs + def allocate( self, candidates: list[tuple[ContextStrategy, float]], total_budget: int, + request: ContextRequest | None = None, ) -> list[tuple[ContextStrategy, float, int]]: """Return (strategy, confidence, allocated_tokens) triples.""" if not candidates: @@ -218,12 +230,12 @@ class ProportionalBudgetAllocator: best = max(candidates, key=lambda x: x[1]) return [(best[0], best[1], total_budget)] - total_confidence = sum(c for _, c in working) - if total_confidence <= 0: - # Equal split when all confidences are zero + total_weight = sum(self._weighted_score(s, c) for s, c in working) + if total_weight <= 0: + # Equal split when all weighted scores are zero return self._equal_split(working, total_budget) - return self._proportional_split(working, total_budget, total_confidence) + return self._proportional_split(working, total_budget, total_weight) def _filter_by_min_budget( self, @@ -231,12 +243,13 @@ class ProportionalBudgetAllocator: total_budget: int, ) -> list[tuple[ContextStrategy, float]]: """Remove candidates whose proportional share < min_useful_budget.""" - total_conf = sum(c for _, c in candidates) - if total_conf <= 0: + total_weight = sum(self._weighted_score(s, c) for s, c in candidates) + if total_weight <= 0: return candidates # Can't compute proportions kept: list[tuple[ContextStrategy, float]] = [] for strategy, confidence in candidates: - share = int(total_budget * confidence / total_conf) + weight = self._weighted_score(strategy, confidence) + share = int(total_budget * weight / total_weight) if share >= self._min_useful_budget: kept.append((strategy, confidence)) if not kept: @@ -261,10 +274,11 @@ class ProportionalBudgetAllocator: self, candidates: list[tuple[ContextStrategy, float]], total_budget: int, - total_confidence: float, + total_weight: float, ) -> list[tuple[ContextStrategy, float, int]]: """Proportional allocation with largest-remainder rounding.""" - raw = [(total_budget * c / total_confidence) for _, c in candidates] + weights = [self._weighted_score(s, c) for s, c in candidates] + raw = [(total_budget * w / total_weight) for w in weights] floors = [int(r) for r in raw] remainder = total_budget - sum(floors) fractions = [ @@ -533,7 +547,7 @@ class ContextAssemblyPipeline(ACMSPipeline): deduplicator=deduplicator or DefaultDeduplicator(), depth_resolver=depth_resolver or DefaultDepthResolver(), scorer=scorer or DefaultScorer(), - packer=packer or DefaultBudgetPacker(), + packer=packer or GreedyKnapsackPacker(), orderer=orderer or DefaultOrderer(), preamble_generator=preamble_generator or DefaultPreambleGenerator(), skeleton_compressor=skeleton_compressor or DefaultSkeletonCompressor(), @@ -552,6 +566,7 @@ class ContextAssemblyPipeline(ACMSPipeline): fragments: Sequence[ContextFragment], budget: ContextBudget, strategy: str | None = None, + request: ContextRequest | None = None, ) -> ContextPayload: """Assemble context with per-stage timing instrumentation. @@ -601,6 +616,7 @@ class ContextAssemblyPipeline(ACMSPipeline): allocations = self._budget_allocator.allocate( candidates, budget.available_tokens, + request, ) allocation_ms = (time.monotonic() - t0) * 1000 @@ -614,7 +630,7 @@ class ContextAssemblyPipeline(ACMSPipeline): dedup_ms = (time.monotonic() - t0) * 1000 t0 = time.monotonic() - fused = self._depth_resolver.resolve(fused) + fused = self._depth_resolver.resolve(fused, budget.available_tokens) depth_ms = (time.monotonic() - t0) * 1000 t0 = time.monotonic() diff --git a/src/cleveragents/application/services/acms_service.py b/src/cleveragents/application/services/acms_service.py index bde322ae3..92dd5b8c3 100644 --- a/src/cleveragents/application/services/acms_service.py +++ b/src/cleveragents/application/services/acms_service.py @@ -33,6 +33,9 @@ if TYPE_CHECKING: from cleveragents.config.settings import Settings from cleveragents.infrastructure.database.unit_of_work import UnitOfWork +from cleveragents.domain.models.acms.crp import ( + ContextRequest, +) from cleveragents.domain.models.core.context_fragment import ( ULID_PATTERN, ContextBudget, @@ -42,6 +45,22 @@ from cleveragents.domain.models.core.context_fragment import ( compute_context_hash, ) +# Lazy import helper — resolved at runtime to avoid circular imports. +_GreedyKnapsackPacker: type | None = None + + +def _get_greedy_knapsack_packer_class() -> type: + """Return the :class:`GreedyKnapsackPacker` class, importing lazily.""" + global _GreedyKnapsackPacker + if _GreedyKnapsackPacker is None: + from cleveragents.application.services.acms_phase2 import ( + GreedyKnapsackPacker, + ) + + _GreedyKnapsackPacker = GreedyKnapsackPacker + return _GreedyKnapsackPacker + + logger = structlog.get_logger(__name__) @@ -58,6 +77,7 @@ class StrategyCapabilities: supports_graph_navigation: bool = False supports_temporal_archaeology: bool = False max_fragments: int | None = None + quality_score: float = 1.0 # --------------------------------------------------------------------------- @@ -260,6 +280,7 @@ class BudgetAllocator(Protocol): self, candidates: list[tuple[ContextStrategy, float]], total_budget: int, + request: ContextRequest | None = None, ) -> list[tuple[ContextStrategy, float, int]]: """Return (strategy, confidence, allocated_tokens) triples.""" ... @@ -302,6 +323,7 @@ class DetailDepthResolver(Protocol): def resolve( self, fragments: Sequence[ContextFragment], + budget: int = 0, ) -> Sequence[ContextFragment]: ... @@ -399,22 +421,31 @@ class DefaultBudgetAllocator: """Distribute the token budget proportionally across candidates. Each candidate receives a share of ``total_budget`` proportional to - its confidence score. When only one candidate is present (v1 norm), - it receives the full budget. With multiple candidates the sum of - allocated tokens never exceeds ``total_budget`` (spec §42689). + ``confidence * quality_score`` per spec §42689. When only one + candidate is present (v1 norm), it receives the full budget. With + multiple candidates the sum of allocated tokens never exceeds + ``total_budget``. """ + @staticmethod + def _weighted_score(strategy: ContextStrategy, confidence: float) -> float: + """Compute the allocation weight: ``confidence * quality_score``.""" + qs = getattr(strategy.capabilities, "quality_score", 1.0) + return confidence * qs + def allocate( self, candidates: list[tuple[ContextStrategy, float]], total_budget: int, + request: ContextRequest | None = None, ) -> list[tuple[ContextStrategy, float, int]]: if not candidates: return [] n = len(candidates) - total_confidence = sum(c for _, c in candidates) - if total_confidence <= 0: - # Equal split when all confidences are zero. + weights = [self._weighted_score(s, c) for s, c in candidates] + total_weight = sum(weights) + if total_weight <= 0: + # Equal split when all weighted scores are zero. share = total_budget // n remainder = total_budget - share * n # Distribute leftover tokens one-per-candidate (largest-remainder). @@ -424,7 +455,7 @@ class DefaultBudgetAllocator: ] # Proportional allocation with largest-remainder distribution so # the full budget is used and no tokens are silently lost. - raw = [(total_budget * c / total_confidence) for _, c in candidates] + raw = [(total_budget * w / total_weight) for w in weights] floors = [int(r) for r in raw] remainder = total_budget - sum(floors) # Award one extra token to candidates with the largest fractional parts. @@ -489,6 +520,7 @@ class DefaultDepthResolver: def resolve( self, fragments: Sequence[ContextFragment], + budget: int = 0, ) -> Sequence[ContextFragment]: return fragments @@ -628,7 +660,7 @@ class ACMSPipeline: self._deduplicator = deduplicator or DefaultDeduplicator() self._depth_resolver = depth_resolver or DefaultDepthResolver() self._scorer = scorer or DefaultScorer() - self._packer = packer or DefaultBudgetPacker() + self._packer = packer or _get_greedy_knapsack_packer_class()() self._orderer = orderer or DefaultOrderer() self._preamble_generator = preamble_generator or DefaultPreambleGenerator() self._skeleton_compressor = skeleton_compressor or DefaultSkeletonCompressor() @@ -639,6 +671,7 @@ class ACMSPipeline: fragments: Sequence[ContextFragment], budget: ContextBudget, strategy: str | None = None, + request: ContextRequest | None = None, ) -> ContextPayload: """Assemble context fragments into a budget-constrained payload.""" if not re.match(ULID_PATTERN, plan_id): @@ -683,12 +716,13 @@ class ACMSPipeline: allocations = self._budget_allocator.allocate( candidates, budget.available_tokens, + request, ) ranked = self._strategy_executor.execute(allocations, fragments, budget) # Phase 2: Fragment Fusion pipeline fused = self._deduplicator.deduplicate(ranked) - fused = self._depth_resolver.resolve(fused) + fused = self._depth_resolver.resolve(fused, budget.available_tokens) fused = self._scorer.score(fused) fused = self._packer.pack(fused, budget) fused = self._orderer.order(fused) diff --git a/src/cleveragents/application/services/fusion_engine.py b/src/cleveragents/application/services/fusion_engine.py index 1353a8bec..86a93992b 100644 --- a/src/cleveragents/application/services/fusion_engine.py +++ b/src/cleveragents/application/services/fusion_engine.py @@ -190,7 +190,7 @@ class FusionEngine: dedup_count = input_count - len(deduped) # Stage 2: Resolve detail depth conflicts (max depth wins) - resolved = list(self._depth_resolver.resolve(deduped)) + resolved = list(self._depth_resolver.resolve(deduped, budget.available_tokens)) depth_resolved_count = len(deduped) - len(resolved) # Stage 3: Score fragments with weighted composite diff --git a/src/cleveragents/application/services/strategy_coordinator.py b/src/cleveragents/application/services/strategy_coordinator.py index 7db72ed8f..c7bceb483 100644 --- a/src/cleveragents/application/services/strategy_coordinator.py +++ b/src/cleveragents/application/services/strategy_coordinator.py @@ -210,7 +210,9 @@ class StrategyCoordinator: return CoordinationResult() # Phase 1b: Budget allocation with per-strategy max cap - allocations = self._allocator.allocate(candidates, budget.available_tokens) + allocations = self._allocator.allocate( + candidates, budget.available_tokens, None + ) capped_allocations = self._apply_max_caps(allocations) # Track circuit-broken strategies