forked from HAL9000/cleveragents-core
feat(acms): implement pipeline orchestrator and Phase 1 components
Add production-quality Phase 1 pipeline components for the ACMS Context Assembly Pipeline: - ConfidenceWeightedSelector: strategy selection with preference boosting and confidence-based ranking - ProportionalBudgetAllocator: proportional token budget distribution with min_useful_budget enforcement and largest-remainder rounding - ParallelStrategyExecutor: concurrent strategy execution via ThreadPoolExecutor with per-strategy timeouts and circuit breaking - CircuitBreaker: per-strategy failure tracking with configurable threshold and explicit reset - ContextAssemblyPipeline: extends ACMSPipeline with Phase 1 production components and per-stage timing (StageTimings) Includes 28 Behave BDD scenarios, 9 Robot Framework integration tests, and ASV benchmarks for all components. ISSUES CLOSED: #539
This commit is contained in:
Forgejo
@@ -0,0 +1,220 @@
|
||||
"""ASV benchmarks for ACMS Pipeline Orchestrator and Phase 1 components.
|
||||
|
||||
Measures the performance of:
|
||||
- ConfidenceWeightedSelector strategy selection
|
||||
- ProportionalBudgetAllocator budget distribution
|
||||
- CircuitBreaker failure tracking
|
||||
- ParallelStrategyExecutor single and multi-strategy execution
|
||||
- ContextAssemblyPipeline full assembly with timing
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
_SRC = str(Path(__file__).resolve().parents[1] / "src")
|
||||
if _SRC not in sys.path:
|
||||
sys.path.insert(0, _SRC)
|
||||
|
||||
import cleveragents # noqa: E402
|
||||
|
||||
importlib.reload(cleveragents)
|
||||
|
||||
from cleveragents.application.services.acms_pipeline import ( # noqa: E402
|
||||
CircuitBreaker,
|
||||
ConfidenceWeightedSelector,
|
||||
ContextAssemblyPipeline,
|
||||
ParallelStrategyExecutor,
|
||||
ProportionalBudgetAllocator,
|
||||
)
|
||||
from cleveragents.application.services.acms_service import ( # noqa: E402
|
||||
RecencyStrategy,
|
||||
RelevanceStrategy,
|
||||
StrategyCapabilities,
|
||||
TieredStrategy,
|
||||
)
|
||||
from cleveragents.domain.models.core.context_fragment import ( # noqa: E402
|
||||
ContextBudget,
|
||||
ContextFragment,
|
||||
FragmentProvenance,
|
||||
)
|
||||
|
||||
_DEFAULT_PROV = FragmentProvenance(resource_uri="bench://orchestrator")
|
||||
|
||||
|
||||
def _make_frag(**kwargs: Any) -> ContextFragment:
|
||||
kwargs.setdefault("uko_node", "bench://orchestrator")
|
||||
kwargs.setdefault("token_count", 10)
|
||||
kwargs.setdefault("provenance", _DEFAULT_PROV)
|
||||
return ContextFragment(**kwargs)
|
||||
|
||||
|
||||
class _BenchStrategy:
|
||||
"""Minimal strategy for benchmarks."""
|
||||
|
||||
def __init__(self, name: str = "bench") -> None:
|
||||
self._name = name
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return self._name
|
||||
|
||||
@property
|
||||
def capabilities(self) -> StrategyCapabilities:
|
||||
return StrategyCapabilities()
|
||||
|
||||
def can_handle(self, request: dict[str, Any]) -> float:
|
||||
return 0.7
|
||||
|
||||
def assemble(
|
||||
self,
|
||||
fragments: Sequence[ContextFragment],
|
||||
budget: ContextBudget,
|
||||
) -> Sequence[ContextFragment]:
|
||||
result: list[ContextFragment] = []
|
||||
total = 0
|
||||
for frag in fragments:
|
||||
if total + frag.token_count <= budget.available_tokens:
|
||||
result.append(frag)
|
||||
total += frag.token_count
|
||||
return result
|
||||
|
||||
def explain(self) -> str:
|
||||
return "Benchmark strategy."
|
||||
|
||||
|
||||
class ConfidenceWeightedSelectorSuite:
|
||||
"""Benchmark ConfidenceWeightedSelector throughput."""
|
||||
|
||||
def setup(self) -> None:
|
||||
self.selector = ConfidenceWeightedSelector()
|
||||
self.strategies = [
|
||||
RelevanceStrategy(),
|
||||
RecencyStrategy(),
|
||||
TieredStrategy(),
|
||||
]
|
||||
|
||||
def time_select_3_strategies(self) -> None:
|
||||
"""Select from 3 built-in strategies."""
|
||||
self.selector.select(self.strategies, {})
|
||||
|
||||
def time_select_with_preferred(self) -> None:
|
||||
"""Select with preferred_strategies boost."""
|
||||
self.selector.select(self.strategies, {"preferred_strategies": ["recency"]})
|
||||
|
||||
|
||||
class ProportionalBudgetAllocatorSuite:
|
||||
"""Benchmark ProportionalBudgetAllocator throughput."""
|
||||
|
||||
def setup(self) -> None:
|
||||
self.allocator = ProportionalBudgetAllocator()
|
||||
self.candidates_2 = [
|
||||
(RelevanceStrategy(), 0.8),
|
||||
(RecencyStrategy(), 0.2),
|
||||
]
|
||||
self.candidates_3 = [
|
||||
(RelevanceStrategy(), 0.5),
|
||||
(RecencyStrategy(), 0.3),
|
||||
(TieredStrategy(), 0.2),
|
||||
]
|
||||
|
||||
def time_allocate_2_candidates(self) -> None:
|
||||
"""Allocate budget across 2 candidates."""
|
||||
self.allocator.allocate(self.candidates_2, 1000)
|
||||
|
||||
def time_allocate_3_candidates(self) -> None:
|
||||
"""Allocate budget across 3 candidates."""
|
||||
self.allocator.allocate(self.candidates_3, 1000)
|
||||
|
||||
|
||||
class CircuitBreakerSuite:
|
||||
"""Benchmark CircuitBreaker operations."""
|
||||
|
||||
def time_record_success(self) -> None:
|
||||
"""Record a success."""
|
||||
cb = CircuitBreaker()
|
||||
cb.record_success("test")
|
||||
|
||||
def time_record_failure(self) -> None:
|
||||
"""Record a failure."""
|
||||
cb = CircuitBreaker()
|
||||
cb.record_failure("test")
|
||||
|
||||
def time_is_open_check(self) -> None:
|
||||
"""Check if circuit is open."""
|
||||
cb = CircuitBreaker()
|
||||
cb.is_open("test")
|
||||
|
||||
|
||||
class ParallelStrategyExecutorSuite:
|
||||
"""Benchmark ParallelStrategyExecutor throughput."""
|
||||
|
||||
def setup(self) -> None:
|
||||
self.executor = ParallelStrategyExecutor()
|
||||
self.fragments = [
|
||||
_make_frag(
|
||||
uko_node=f"bench://file{i}.py",
|
||||
content=f"content_{i}",
|
||||
token_count=10,
|
||||
relevance_score=0.5,
|
||||
)
|
||||
for i in range(20)
|
||||
]
|
||||
self.budget = ContextBudget(max_tokens=500, reserved_tokens=0)
|
||||
self.strategy = _BenchStrategy()
|
||||
|
||||
def time_execute_single_strategy(self) -> None:
|
||||
"""Execute a single strategy (synchronous path)."""
|
||||
self.executor.execute(
|
||||
[(self.strategy, 0.7, 200)],
|
||||
self.fragments,
|
||||
self.budget,
|
||||
)
|
||||
|
||||
|
||||
class ContextAssemblyPipelineSuite:
|
||||
"""Benchmark full pipeline assembly."""
|
||||
|
||||
def setup(self) -> None:
|
||||
self.pipeline = ContextAssemblyPipeline()
|
||||
self.budget = ContextBudget(max_tokens=500, reserved_tokens=0)
|
||||
self.fragments_small = [
|
||||
_make_frag(
|
||||
uko_node=f"bench://s{i}.py",
|
||||
content=f"small_{i}",
|
||||
token_count=10,
|
||||
relevance_score=round(0.9 - i * 0.05, 2),
|
||||
)
|
||||
for i in range(5)
|
||||
]
|
||||
self.fragments_medium = [
|
||||
_make_frag(
|
||||
uko_node=f"bench://m{i}.py",
|
||||
content=f"medium_{i}",
|
||||
token_count=10,
|
||||
relevance_score=round(0.9 - i * 0.01, 2),
|
||||
)
|
||||
for i in range(50)
|
||||
]
|
||||
|
||||
def time_assemble_5_fragments(self) -> None:
|
||||
"""Assemble 5 fragments via full pipeline."""
|
||||
self.pipeline.assemble(
|
||||
plan_id="01JQTESTPN00000000000000AA",
|
||||
fragments=self.fragments_small,
|
||||
budget=self.budget,
|
||||
strategy="relevance",
|
||||
)
|
||||
|
||||
def time_assemble_50_fragments(self) -> None:
|
||||
"""Assemble 50 fragments via full pipeline."""
|
||||
self.pipeline.assemble(
|
||||
plan_id="01JQTESTPN00000000000000AA",
|
||||
fragments=self.fragments_medium,
|
||||
budget=self.budget,
|
||||
strategy="relevance",
|
||||
)
|
||||
@@ -0,0 +1,267 @@
|
||||
@phase2 @acms @acms_pipeline_orchestrator
|
||||
Feature: ACMS Pipeline Orchestrator and Phase 1 Components
|
||||
As a CleverAgents developer
|
||||
I want production-quality Phase 1 pipeline components
|
||||
So that context strategies are selected, budgeted, and executed
|
||||
with confidence weighting, proportional allocation, parallel
|
||||
execution, and circuit-breaker fault tolerance
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ConfidenceWeightedSelector
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@orchestrator @selector
|
||||
Scenario: ConfidenceWeightedSelector selects strategies by confidence
|
||||
Given the pipeline orchestrator modules are available
|
||||
When I use ConfidenceWeightedSelector to select from 3 strategies
|
||||
Then all strategies with positive confidence should be returned
|
||||
And the results should be sorted by confidence descending
|
||||
|
||||
@orchestrator @selector
|
||||
Scenario: ConfidenceWeightedSelector boosts preferred strategies
|
||||
Given the pipeline orchestrator modules are available
|
||||
When I select with preferred_strategies containing "recency"
|
||||
Then the "recency" strategy should have a boosted confidence
|
||||
And the boosted confidence should not exceed 1.0
|
||||
|
||||
@orchestrator @selector
|
||||
Scenario: ConfidenceWeightedSelector excludes zero-confidence strategies
|
||||
Given the pipeline orchestrator modules are available
|
||||
And a test strategy with zero confidence
|
||||
When I use ConfidenceWeightedSelector with the zero-confidence strategy
|
||||
Then the zero-confidence strategy should not be in the results
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ProportionalBudgetAllocator
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@orchestrator @allocator
|
||||
Scenario: ProportionalBudgetAllocator distributes proportionally
|
||||
Given the pipeline orchestrator modules are available
|
||||
When I use ProportionalBudgetAllocator for candidates with confidences 0.8 and 0.2 and budget 1000
|
||||
Then the first candidate should receive approximately 800 tokens from proportional allocator
|
||||
And the second candidate should receive approximately 200 tokens from proportional allocator
|
||||
And the total proportional allocation should equal exactly 1000
|
||||
|
||||
@orchestrator @allocator
|
||||
Scenario: ProportionalBudgetAllocator enforces min_useful_budget
|
||||
Given the pipeline orchestrator modules are available
|
||||
When I use ProportionalBudgetAllocator with min_useful_budget 200 for candidates with confidences 0.99 and 0.01 and budget 300
|
||||
Then only the high-confidence candidate should receive tokens
|
||||
And the total proportional allocation should equal exactly 300
|
||||
|
||||
@orchestrator @allocator
|
||||
Scenario: ProportionalBudgetAllocator handles single candidate
|
||||
Given the pipeline orchestrator modules are available
|
||||
When I use ProportionalBudgetAllocator for a single candidate with budget 500
|
||||
Then the single candidate should receive all 500 tokens
|
||||
|
||||
@orchestrator @allocator
|
||||
Scenario: ProportionalBudgetAllocator handles zero confidence with equal split
|
||||
Given the pipeline orchestrator modules are available
|
||||
When I use ProportionalBudgetAllocator for 3 zero-confidence candidates with budget 999
|
||||
Then the total proportional allocation should equal exactly 999
|
||||
And each proportional allocation should be 333 or 334
|
||||
|
||||
@orchestrator @allocator
|
||||
Scenario: ProportionalBudgetAllocator falls back when all excluded
|
||||
Given the pipeline orchestrator modules are available
|
||||
When I use ProportionalBudgetAllocator with min_useful_budget 1000 for 2 candidates with budget 500
|
||||
Then the highest-confidence candidate should receive the full budget
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CircuitBreaker
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@orchestrator @circuit_breaker
|
||||
Scenario: CircuitBreaker opens after threshold failures
|
||||
Given the pipeline orchestrator modules are available
|
||||
And a circuit breaker with threshold 3
|
||||
When I record 3 consecutive failures for strategy "flaky"
|
||||
Then the circuit for "flaky" should be open
|
||||
|
||||
@orchestrator @circuit_breaker
|
||||
Scenario: CircuitBreaker resets on success
|
||||
Given the pipeline orchestrator modules are available
|
||||
And a circuit breaker with threshold 3
|
||||
When I record 2 failures then 1 success for strategy "recovering"
|
||||
Then the circuit for "recovering" should be closed
|
||||
|
||||
@orchestrator @circuit_breaker
|
||||
Scenario: CircuitBreaker can be explicitly reset
|
||||
Given the pipeline orchestrator modules are available
|
||||
And a circuit breaker with threshold 2
|
||||
When I record 2 failures for strategy "resettable"
|
||||
And I reset the circuit for "resettable"
|
||||
Then the circuit for "resettable" should be closed
|
||||
|
||||
@orchestrator @circuit_breaker
|
||||
Scenario: CircuitBreaker reset_all clears all circuits
|
||||
Given the pipeline orchestrator modules are available
|
||||
And a circuit breaker with threshold 1
|
||||
When I record 1 failure for strategy "a" and 1 failure for strategy "b"
|
||||
And I reset all circuits
|
||||
Then the circuit for "a" should be closed
|
||||
And the circuit for "b" should be closed
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ParallelStrategyExecutor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@orchestrator @executor
|
||||
Scenario: ParallelStrategyExecutor invokes strategy and returns fragments
|
||||
Given the pipeline orchestrator modules are available
|
||||
And a tracking test strategy
|
||||
And context fragments for executor testing
|
||||
When I execute the strategy via ParallelStrategyExecutor
|
||||
Then the executor should return fragments from the strategy
|
||||
And the tracking strategy should have been invoked
|
||||
|
||||
@orchestrator @executor
|
||||
Scenario: ParallelStrategyExecutor skips circuit-broken strategies
|
||||
Given the pipeline orchestrator modules are available
|
||||
And a ParallelStrategyExecutor with a pre-broken circuit for "broken-strat"
|
||||
And context fragments for executor testing
|
||||
When I execute with the circuit-broken strategy
|
||||
Then the executor should return 0 fragments
|
||||
|
||||
@orchestrator @executor
|
||||
Scenario: ParallelStrategyExecutor handles strategy failures gracefully
|
||||
Given the pipeline orchestrator modules are available
|
||||
And a strategy that always raises an exception
|
||||
And context fragments for executor testing
|
||||
When I execute the failing strategy via ParallelStrategyExecutor
|
||||
Then the executor should return 0 fragments
|
||||
And the circuit breaker should have recorded a failure for "failing"
|
||||
|
||||
@orchestrator @executor
|
||||
Scenario: ParallelStrategyExecutor skips zero-budget allocations
|
||||
Given the pipeline orchestrator modules are available
|
||||
And a tracking test strategy
|
||||
And context fragments for executor testing
|
||||
When I execute with zero-budget allocation
|
||||
Then the executor should return 0 fragments
|
||||
|
||||
@orchestrator @executor
|
||||
Scenario: ParallelStrategyExecutor executes multiple strategies in parallel
|
||||
Given the pipeline orchestrator modules are available
|
||||
And context fragments for executor testing
|
||||
When I execute 2 tracking strategies in parallel via ParallelStrategyExecutor
|
||||
Then the executor should return fragments from all parallel strategies
|
||||
And both parallel strategies should have been invoked
|
||||
|
||||
@orchestrator @executor
|
||||
Scenario: ParallelStrategyExecutor handles mixed success and failure in parallel
|
||||
Given the pipeline orchestrator modules are available
|
||||
And context fragments for executor testing
|
||||
When I execute a tracking strategy and a failing strategy in parallel
|
||||
Then the executor should return fragments only from the successful parallel strategy
|
||||
And the parallel circuit breaker should record a failure for "failing"
|
||||
|
||||
@orchestrator @executor
|
||||
Scenario: ParallelStrategyExecutor returns empty for empty allocations
|
||||
Given the pipeline orchestrator modules are available
|
||||
And context fragments for executor testing
|
||||
When I execute with empty allocations via ParallelStrategyExecutor
|
||||
Then the executor should return 0 fragments
|
||||
|
||||
@orchestrator @executor
|
||||
Scenario: ParallelStrategyExecutor circuit_breaker property is accessible
|
||||
Given the pipeline orchestrator modules are available
|
||||
When I create a ParallelStrategyExecutor with a custom circuit breaker
|
||||
Then the circuit_breaker property should return the custom breaker
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ProportionalBudgetAllocator — Additional edge cases
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@orchestrator @allocator
|
||||
Scenario: ProportionalBudgetAllocator handles empty candidate list
|
||||
Given the pipeline orchestrator modules are available
|
||||
When I use ProportionalBudgetAllocator for empty candidates with budget 1000
|
||||
Then the proportional allocator should return an empty list
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ContextAssemblyPipeline — Additional edge cases
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@orchestrator @pipeline
|
||||
Scenario: ContextAssemblyPipeline rejects unknown strategy name
|
||||
Given a ContextAssemblyPipeline with default components
|
||||
And the following orchestrator test fragments:
|
||||
| uko_node | content | score | tokens |
|
||||
| project://app/a.py | alpha | 0.9 | 100 |
|
||||
And a pipeline budget with max_tokens 500 and reserved_tokens 0
|
||||
When I assemble via the orchestrator with strategy "nonexistent_strategy"
|
||||
Then an orchestrator ValueError should be raised mentioning "Unknown strategy"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ContextAssemblyPipeline — Full integration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@orchestrator @pipeline
|
||||
Scenario: ContextAssemblyPipeline assembles with timing metadata
|
||||
Given a ContextAssemblyPipeline with default components
|
||||
And the following orchestrator test fragments:
|
||||
| uko_node | content | score | tokens |
|
||||
| project://app/a.py | alpha | 0.9 | 100 |
|
||||
| project://app/b.py | beta | 0.3 | 100 |
|
||||
And a pipeline budget with max_tokens 250 and reserved_tokens 0
|
||||
When I assemble via the orchestrator with strategy "relevance"
|
||||
Then the orchestrator payload should contain 2 fragments
|
||||
And the orchestrator should have timing metadata
|
||||
And all timing values should be non-negative
|
||||
|
||||
@orchestrator @pipeline
|
||||
Scenario: ContextAssemblyPipeline inherits ACMSPipeline behavior
|
||||
Given a ContextAssemblyPipeline with default components
|
||||
And the following orchestrator test fragments:
|
||||
| uko_node | content | score | tokens |
|
||||
| project://app/a.py | alpha | 0.9 | 50 |
|
||||
And a pipeline budget with max_tokens 500 and reserved_tokens 0
|
||||
When I assemble via the orchestrator with strategy "relevance"
|
||||
Then the orchestrator payload strategies used should include "relevance"
|
||||
And the orchestrator payload should be within budget
|
||||
|
||||
@orchestrator @pipeline
|
||||
Scenario: ContextAssemblyPipeline rejects invalid plan_id
|
||||
Given a ContextAssemblyPipeline with default components
|
||||
And the following orchestrator test fragments:
|
||||
| uko_node | content | score | tokens |
|
||||
| project://app/a.py | alpha | 0.9 | 100 |
|
||||
And a pipeline budget with max_tokens 500 and reserved_tokens 0
|
||||
When I assemble via the orchestrator with invalid plan_id "not-valid"
|
||||
Then an orchestrator ValueError should be raised mentioning "ULID"
|
||||
|
||||
@orchestrator @pipeline
|
||||
Scenario: ContextAssemblyPipeline accepts custom Phase 1 components
|
||||
Given a ContextAssemblyPipeline with a custom tracking selector
|
||||
And the following orchestrator test fragments:
|
||||
| uko_node | content | score | tokens |
|
||||
| project://app/a.py | alpha | 0.9 | 100 |
|
||||
And a pipeline budget with max_tokens 500 and reserved_tokens 0
|
||||
When I assemble via the orchestrator with strategy "relevance"
|
||||
Then the custom orchestrator selector should have been called
|
||||
|
||||
@orchestrator @pipeline
|
||||
Scenario: ContextAssemblyPipeline can register custom strategies
|
||||
Given a ContextAssemblyPipeline with default 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 |
|
||||
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"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# StageTimings data model
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@orchestrator @timings
|
||||
Scenario: StageTimings is a frozen dataclass with named fields
|
||||
Given the pipeline orchestrator modules are available
|
||||
When I create a StageTimings with total_ms 42.5
|
||||
Then the timings total_ms should be 42.5
|
||||
And the timings should have all 10 named fields
|
||||
@@ -0,0 +1,854 @@
|
||||
"""Step definitions for features/acms_pipeline_orchestrator.feature.
|
||||
|
||||
Tests the ACMS Pipeline Orchestrator and Phase 1 production components:
|
||||
ConfidenceWeightedSelector, ProportionalBudgetAllocator,
|
||||
ParallelStrategyExecutor, CircuitBreaker, ContextAssemblyPipeline,
|
||||
and StageTimings.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
from behave import given, then, when
|
||||
from behave.runner import Context
|
||||
|
||||
from cleveragents.application.services.acms_pipeline import (
|
||||
CircuitBreaker,
|
||||
ConfidenceWeightedSelector,
|
||||
ContextAssemblyPipeline,
|
||||
ParallelStrategyExecutor,
|
||||
ProportionalBudgetAllocator,
|
||||
StageTimings,
|
||||
)
|
||||
from cleveragents.application.services.acms_service import (
|
||||
DefaultStrategySelector,
|
||||
RecencyStrategy,
|
||||
RelevanceStrategy,
|
||||
StrategyCapabilities,
|
||||
TieredStrategy,
|
||||
)
|
||||
from cleveragents.domain.models.core.context_fragment import (
|
||||
ContextBudget,
|
||||
ContextFragment,
|
||||
FragmentProvenance,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_DEFAULT_PROVENANCE = FragmentProvenance(resource_uri="test://orchestrator")
|
||||
|
||||
|
||||
def _make_frag(**kwargs: Any) -> ContextFragment:
|
||||
"""Create a ContextFragment with test defaults."""
|
||||
kwargs.setdefault("uko_node", "test://orchestrator")
|
||||
kwargs.setdefault("token_count", 0)
|
||||
kwargs.setdefault("provenance", _DEFAULT_PROVENANCE)
|
||||
return ContextFragment(**kwargs)
|
||||
|
||||
|
||||
class _ZeroConfidenceStrategy:
|
||||
"""Test strategy that always returns 0.0 confidence."""
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "zero_conf"
|
||||
|
||||
@property
|
||||
def capabilities(self) -> StrategyCapabilities:
|
||||
return StrategyCapabilities()
|
||||
|
||||
def can_handle(self, request: dict[str, Any]) -> float:
|
||||
return 0.0
|
||||
|
||||
def assemble(
|
||||
self,
|
||||
fragments: Sequence[ContextFragment],
|
||||
budget: ContextBudget,
|
||||
) -> Sequence[ContextFragment]:
|
||||
return []
|
||||
|
||||
def explain(self) -> str:
|
||||
return "Always zero confidence."
|
||||
|
||||
|
||||
class _TrackingTestStrategy:
|
||||
"""Test strategy that tracks invocation and returns all fragments within budget."""
|
||||
|
||||
def __init__(self, name: str = "tracking") -> None:
|
||||
self._name = name
|
||||
self.invoked = False
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return self._name
|
||||
|
||||
@property
|
||||
def capabilities(self) -> StrategyCapabilities:
|
||||
return StrategyCapabilities()
|
||||
|
||||
def can_handle(self, request: dict[str, Any]) -> float:
|
||||
return 0.7
|
||||
|
||||
def assemble(
|
||||
self,
|
||||
fragments: Sequence[ContextFragment],
|
||||
budget: ContextBudget,
|
||||
) -> Sequence[ContextFragment]:
|
||||
self.invoked = True
|
||||
result: list[ContextFragment] = []
|
||||
total = 0
|
||||
for frag in fragments:
|
||||
if total + frag.token_count <= budget.available_tokens:
|
||||
result.append(frag)
|
||||
total += frag.token_count
|
||||
return result
|
||||
|
||||
def explain(self) -> str:
|
||||
return "Tracking strategy for testing."
|
||||
|
||||
|
||||
class _FailingStrategy:
|
||||
"""Test strategy that always raises an exception."""
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "failing"
|
||||
|
||||
@property
|
||||
def capabilities(self) -> StrategyCapabilities:
|
||||
return StrategyCapabilities()
|
||||
|
||||
def can_handle(self, request: dict[str, Any]) -> float:
|
||||
return 0.5
|
||||
|
||||
def assemble(
|
||||
self,
|
||||
fragments: Sequence[ContextFragment],
|
||||
budget: ContextBudget,
|
||||
) -> Sequence[ContextFragment]:
|
||||
msg = "Intentional test failure"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
def explain(self) -> str:
|
||||
return "Always fails."
|
||||
|
||||
|
||||
class _ReverseOrchestratorStrategy:
|
||||
"""Test strategy that reverses fragment order."""
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "reverse_orch"
|
||||
|
||||
@property
|
||||
def capabilities(self) -> StrategyCapabilities:
|
||||
return StrategyCapabilities()
|
||||
|
||||
def can_handle(self, request: dict[str, Any]) -> float:
|
||||
return 0.5
|
||||
|
||||
def assemble(
|
||||
self,
|
||||
fragments: Sequence[ContextFragment],
|
||||
budget: ContextBudget,
|
||||
) -> Sequence[ContextFragment]:
|
||||
reversed_frags = list(reversed(fragments))
|
||||
result: list[ContextFragment] = []
|
||||
total = 0
|
||||
for frag in reversed_frags:
|
||||
if total + frag.token_count <= budget.available_tokens:
|
||||
result.append(frag)
|
||||
total += frag.token_count
|
||||
return result
|
||||
|
||||
def explain(self) -> str:
|
||||
return "Reverses fragment order."
|
||||
|
||||
|
||||
class _TrackingOrchestratorSelector:
|
||||
"""Strategy selector that tracks invocation and delegates to default."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.called = False
|
||||
|
||||
def select(
|
||||
self,
|
||||
strategies: Sequence[Any],
|
||||
request: dict[str, Any],
|
||||
) -> list[tuple[Any, float]]:
|
||||
self.called = True
|
||||
return DefaultStrategySelector().select(strategies, request)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Given steps
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@given("the pipeline orchestrator modules are available")
|
||||
def step_orchestrator_modules_available(context: Context) -> None:
|
||||
"""Ensure all pipeline orchestrator modules are importable."""
|
||||
pass
|
||||
|
||||
|
||||
@given("a test strategy with zero confidence")
|
||||
def step_zero_confidence_strategy(context: Context) -> None:
|
||||
context.zero_conf_strategy = _ZeroConfidenceStrategy()
|
||||
|
||||
|
||||
@given("a circuit breaker with threshold {threshold:d}")
|
||||
def step_circuit_breaker(context: Context, threshold: int) -> None:
|
||||
context.circuit_breaker = CircuitBreaker(failure_threshold=threshold)
|
||||
|
||||
|
||||
@given("a tracking test strategy")
|
||||
def step_tracking_strategy(context: Context) -> None:
|
||||
context.tracking_strategy = _TrackingTestStrategy()
|
||||
|
||||
|
||||
@given("context fragments for executor testing")
|
||||
def step_executor_fragments(context: Context) -> None:
|
||||
context.executor_fragments = [
|
||||
_make_frag(
|
||||
uko_node="project://test/a.py",
|
||||
content="alpha",
|
||||
token_count=50,
|
||||
relevance_score=0.9,
|
||||
),
|
||||
_make_frag(
|
||||
uko_node="project://test/b.py",
|
||||
content="beta",
|
||||
token_count=50,
|
||||
relevance_score=0.5,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@given('a ParallelStrategyExecutor with a pre-broken circuit for "{name}"')
|
||||
def step_pre_broken_executor(context: Context, name: str) -> None:
|
||||
cb = CircuitBreaker(failure_threshold=1)
|
||||
cb.record_failure(name)
|
||||
context.broken_executor = ParallelStrategyExecutor(circuit_breaker=cb)
|
||||
context.broken_strategy_name = name
|
||||
|
||||
|
||||
@given("a strategy that always raises an exception")
|
||||
def step_failing_strategy(context: Context) -> None:
|
||||
context.failing_strategy = _FailingStrategy()
|
||||
|
||||
|
||||
@given("a ContextAssemblyPipeline with default components")
|
||||
def step_default_pipeline(context: Context) -> None:
|
||||
context.orch_pipeline = ContextAssemblyPipeline()
|
||||
|
||||
|
||||
@given("the following orchestrator test fragments:")
|
||||
def step_orch_fragments_table(context: Context) -> None:
|
||||
context.orch_fragments = []
|
||||
for row in context.table:
|
||||
uko_node = row["uko_node"]
|
||||
frag = ContextFragment(
|
||||
uko_node=uko_node,
|
||||
content=row["content"],
|
||||
relevance_score=float(row["score"]),
|
||||
token_count=int(row["tokens"]),
|
||||
provenance=FragmentProvenance(resource_uri=uko_node),
|
||||
)
|
||||
context.orch_fragments.append(frag)
|
||||
|
||||
|
||||
@given("a pipeline budget with max_tokens {max_tok:d} and reserved_tokens {res:d}")
|
||||
def step_pipeline_budget(context: Context, max_tok: int, res: int) -> None:
|
||||
context.orch_budget = ContextBudget(max_tokens=max_tok, reserved_tokens=res)
|
||||
|
||||
|
||||
@given("a ContextAssemblyPipeline with a custom tracking selector")
|
||||
def step_pipeline_custom_selector(context: Context) -> None:
|
||||
selector = _TrackingOrchestratorSelector()
|
||||
context.orch_custom_selector = selector
|
||||
context.orch_pipeline = ContextAssemblyPipeline(strategy_selector=selector)
|
||||
|
||||
|
||||
@given("a custom orchestrator strategy that returns fragments in reverse")
|
||||
def step_custom_reverse_strategy(context: Context) -> None:
|
||||
if not hasattr(context, "orch_pipeline"):
|
||||
context.orch_pipeline = ContextAssemblyPipeline()
|
||||
context.orch_pipeline.register_strategy(
|
||||
"reverse_orch", _ReverseOrchestratorStrategy()
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# When steps — ConfidenceWeightedSelector
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@when("I use ConfidenceWeightedSelector to select from 3 strategies")
|
||||
def step_cws_select(context: Context) -> None:
|
||||
strategies = [RelevanceStrategy(), RecencyStrategy(), TieredStrategy()]
|
||||
selector = ConfidenceWeightedSelector()
|
||||
context.cws_results = selector.select(strategies, {})
|
||||
|
||||
|
||||
@when('I select with preferred_strategies containing "{preferred}"')
|
||||
def step_cws_preferred(context: Context, preferred: str) -> None:
|
||||
strategies = [RelevanceStrategy(), RecencyStrategy(), TieredStrategy()]
|
||||
selector = ConfidenceWeightedSelector(preference_boost=1.5)
|
||||
context.cws_results_preferred = selector.select(
|
||||
strategies, {"preferred_strategies": [preferred]}
|
||||
)
|
||||
# Also get non-boosted results for comparison
|
||||
context.cws_results_baseline = selector.select(strategies, {})
|
||||
context.preferred_name = preferred
|
||||
|
||||
|
||||
@when("I use ConfidenceWeightedSelector with the zero-confidence strategy")
|
||||
def step_cws_zero_conf(context: Context) -> None:
|
||||
strategies = [RelevanceStrategy(), context.zero_conf_strategy]
|
||||
selector = ConfidenceWeightedSelector()
|
||||
context.cws_results_with_zero = selector.select(strategies, {})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# When steps — ProportionalBudgetAllocator
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@when(
|
||||
"I use ProportionalBudgetAllocator for candidates with confidences {c1:g} and {c2:g} and budget {budget:d}"
|
||||
)
|
||||
def step_pba_proportional(context: Context, c1: float, c2: float, budget: int) -> None:
|
||||
s1 = RelevanceStrategy()
|
||||
s2 = RecencyStrategy()
|
||||
allocator = ProportionalBudgetAllocator()
|
||||
context.pba_allocations = allocator.allocate([(s1, c1), (s2, c2)], budget)
|
||||
context.pba_budget = budget
|
||||
|
||||
|
||||
@when(
|
||||
"I use ProportionalBudgetAllocator with min_useful_budget {min_budget:d} for candidates with confidences {c1:g} and {c2:g} and budget {budget:d}"
|
||||
)
|
||||
def step_pba_min_budget(
|
||||
context: Context, min_budget: int, c1: float, c2: float, budget: int
|
||||
) -> None:
|
||||
s1 = RelevanceStrategy()
|
||||
s2 = RecencyStrategy()
|
||||
allocator = ProportionalBudgetAllocator(min_useful_budget=min_budget)
|
||||
context.pba_allocations = allocator.allocate([(s1, c1), (s2, c2)], budget)
|
||||
context.pba_budget = budget
|
||||
|
||||
|
||||
@when("I use ProportionalBudgetAllocator for a single candidate with budget {budget:d}")
|
||||
def step_pba_single(context: Context, budget: int) -> None:
|
||||
s1 = RelevanceStrategy()
|
||||
allocator = ProportionalBudgetAllocator()
|
||||
context.pba_allocations = allocator.allocate([(s1, 0.8)], budget)
|
||||
context.pba_budget = budget
|
||||
|
||||
|
||||
@when(
|
||||
"I use ProportionalBudgetAllocator for {count:d} zero-confidence candidates with budget {budget:d}"
|
||||
)
|
||||
def step_pba_zero_conf(context: Context, count: int, budget: int) -> None:
|
||||
strategies = [RelevanceStrategy(), RecencyStrategy(), TieredStrategy()]
|
||||
allocator = ProportionalBudgetAllocator()
|
||||
candidates = [(s, 0.0) for s in strategies[:count]]
|
||||
context.pba_allocations = allocator.allocate(candidates, budget)
|
||||
context.pba_budget = budget
|
||||
|
||||
|
||||
@when(
|
||||
"I use ProportionalBudgetAllocator with min_useful_budget {min_budget:d} for {count:d} candidates with budget {budget:d}"
|
||||
)
|
||||
def step_pba_all_excluded(
|
||||
context: Context, min_budget: int, count: int, budget: int
|
||||
) -> None:
|
||||
strategies = [RelevanceStrategy(), RecencyStrategy()]
|
||||
allocator = ProportionalBudgetAllocator(min_useful_budget=min_budget)
|
||||
candidates = [(s, 0.5) for s in strategies[:count]]
|
||||
context.pba_allocations = allocator.allocate(candidates, budget)
|
||||
context.pba_budget = budget
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# When steps — CircuitBreaker
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@when('I record {count:d} consecutive failures for strategy "{name}"')
|
||||
def step_cb_failures(context: Context, count: int, name: str) -> None:
|
||||
for _ in range(count):
|
||||
context.circuit_breaker.record_failure(name)
|
||||
|
||||
|
||||
@when('I record 2 failures then 1 success for strategy "{name}"')
|
||||
def step_cb_failures_then_success(context: Context, name: str) -> None:
|
||||
context.circuit_breaker.record_failure(name)
|
||||
context.circuit_breaker.record_failure(name)
|
||||
context.circuit_breaker.record_success(name)
|
||||
|
||||
|
||||
@when('I record 2 failures for strategy "{name}"')
|
||||
def step_cb_2_failures(context: Context, name: str) -> None:
|
||||
context.circuit_breaker.record_failure(name)
|
||||
context.circuit_breaker.record_failure(name)
|
||||
|
||||
|
||||
@when('I reset the circuit for "{name}"')
|
||||
def step_cb_reset(context: Context, name: str) -> None:
|
||||
context.circuit_breaker.reset(name)
|
||||
|
||||
|
||||
@when('I record 1 failure for strategy "{name1}" and 1 failure for strategy "{name2}"')
|
||||
def step_cb_two_strat_failures(context: Context, name1: str, name2: str) -> None:
|
||||
context.circuit_breaker.record_failure(name1)
|
||||
context.circuit_breaker.record_failure(name2)
|
||||
|
||||
|
||||
@when("I reset all circuits")
|
||||
def step_cb_reset_all(context: Context) -> None:
|
||||
context.circuit_breaker.reset_all()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# When steps — ParallelStrategyExecutor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@when("I execute the strategy via ParallelStrategyExecutor")
|
||||
def step_executor_run(context: Context) -> None:
|
||||
executor = ParallelStrategyExecutor()
|
||||
budget = ContextBudget(max_tokens=200, reserved_tokens=0)
|
||||
allocations = [(context.tracking_strategy, 0.7, 200)]
|
||||
context.executor_results = executor.execute(
|
||||
allocations, context.executor_fragments, budget
|
||||
)
|
||||
|
||||
|
||||
@when("I execute with the circuit-broken strategy")
|
||||
def step_executor_broken(context: Context) -> None:
|
||||
strategy = _TrackingTestStrategy(name=context.broken_strategy_name)
|
||||
budget = ContextBudget(max_tokens=200, reserved_tokens=0)
|
||||
allocations = [(strategy, 0.7, 200)]
|
||||
context.executor_results = context.broken_executor.execute(
|
||||
allocations, context.executor_fragments, budget
|
||||
)
|
||||
|
||||
|
||||
@when("I execute the failing strategy via ParallelStrategyExecutor")
|
||||
def step_executor_failing(context: Context) -> None:
|
||||
cb = CircuitBreaker()
|
||||
executor = ParallelStrategyExecutor(circuit_breaker=cb)
|
||||
budget = ContextBudget(max_tokens=200, reserved_tokens=0)
|
||||
allocations = [(context.failing_strategy, 0.5, 200)]
|
||||
context.executor_results = executor.execute(
|
||||
allocations, context.executor_fragments, budget
|
||||
)
|
||||
context.executor_cb = cb
|
||||
|
||||
|
||||
@when("I execute with zero-budget allocation")
|
||||
def step_executor_zero_budget(context: Context) -> None:
|
||||
executor = ParallelStrategyExecutor()
|
||||
budget = ContextBudget(max_tokens=200, reserved_tokens=0)
|
||||
allocations = [(context.tracking_strategy, 0.7, 0)]
|
||||
context.executor_results = executor.execute(
|
||||
allocations, context.executor_fragments, budget
|
||||
)
|
||||
|
||||
|
||||
@when("I execute 2 tracking strategies in parallel via ParallelStrategyExecutor")
|
||||
def step_executor_parallel(context: Context) -> None:
|
||||
s1 = _TrackingTestStrategy(name="parallel_a")
|
||||
s2 = _TrackingTestStrategy(name="parallel_b")
|
||||
context.parallel_strategies = [s1, s2]
|
||||
executor = ParallelStrategyExecutor(max_workers=2)
|
||||
budget = ContextBudget(max_tokens=200, reserved_tokens=0)
|
||||
allocations: list[tuple[Any, float, int]] = [
|
||||
(s1, 0.6, 100),
|
||||
(s2, 0.4, 100),
|
||||
]
|
||||
context.executor_results = executor.execute(
|
||||
allocations, context.executor_fragments, budget
|
||||
)
|
||||
|
||||
|
||||
@when("I execute a tracking strategy and a failing strategy in parallel")
|
||||
def step_executor_parallel_mixed(context: Context) -> None:
|
||||
s_ok = _TrackingTestStrategy(name="good_parallel")
|
||||
s_fail = _FailingStrategy()
|
||||
context.parallel_ok_strategy = s_ok
|
||||
cb = CircuitBreaker()
|
||||
executor = ParallelStrategyExecutor(max_workers=2, circuit_breaker=cb)
|
||||
budget = ContextBudget(max_tokens=200, reserved_tokens=0)
|
||||
allocations: list[tuple[Any, float, int]] = [
|
||||
(s_ok, 0.6, 100),
|
||||
(s_fail, 0.4, 100),
|
||||
]
|
||||
context.executor_results = executor.execute(
|
||||
allocations, context.executor_fragments, budget
|
||||
)
|
||||
context.parallel_cb = cb
|
||||
|
||||
|
||||
@when("I execute with empty allocations via ParallelStrategyExecutor")
|
||||
def step_executor_empty_allocations(context: Context) -> None:
|
||||
executor = ParallelStrategyExecutor()
|
||||
budget = ContextBudget(max_tokens=200, reserved_tokens=0)
|
||||
context.executor_results = executor.execute([], context.executor_fragments, budget)
|
||||
|
||||
|
||||
@when("I create a ParallelStrategyExecutor with a custom circuit breaker")
|
||||
def step_executor_custom_cb(context: Context) -> None:
|
||||
custom_cb = CircuitBreaker(failure_threshold=10)
|
||||
context.custom_executor = ParallelStrategyExecutor(circuit_breaker=custom_cb)
|
||||
context.custom_cb = custom_cb
|
||||
|
||||
|
||||
@when("I use ProportionalBudgetAllocator for empty candidates with budget {budget:d}")
|
||||
def step_pba_empty_candidates(context: Context, budget: int) -> None:
|
||||
allocator = ProportionalBudgetAllocator()
|
||||
context.pba_allocations = allocator.allocate([], budget)
|
||||
context.pba_budget = budget
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# When steps — ContextAssemblyPipeline
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@when('I assemble via the orchestrator with strategy "{strategy}"')
|
||||
def step_orch_assemble(context: Context, strategy: str) -> None:
|
||||
context.orch_error = None
|
||||
try:
|
||||
context.orch_payload = context.orch_pipeline.assemble(
|
||||
plan_id="01JQTESTPN00000000000000AA",
|
||||
fragments=list(context.orch_fragments),
|
||||
budget=context.orch_budget,
|
||||
strategy=strategy,
|
||||
)
|
||||
except ValueError as exc:
|
||||
context.orch_error = exc
|
||||
|
||||
|
||||
@when('I assemble via the orchestrator with invalid plan_id "{plan_id}"')
|
||||
def step_orch_assemble_invalid(context: Context, plan_id: str) -> None:
|
||||
context.orch_error = None
|
||||
try:
|
||||
context.orch_payload = context.orch_pipeline.assemble(
|
||||
plan_id=plan_id,
|
||||
fragments=list(context.orch_fragments),
|
||||
budget=context.orch_budget,
|
||||
)
|
||||
except ValueError as exc:
|
||||
context.orch_error = exc
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# When steps — StageTimings
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@when("I create a StageTimings with total_ms {total:g}")
|
||||
def step_create_timings(context: Context, total: float) -> None:
|
||||
context.test_timings = StageTimings(total_ms=total)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Then steps — ConfidenceWeightedSelector
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@then("all strategies with positive confidence should be returned")
|
||||
def step_cws_all_positive(context: Context) -> None:
|
||||
assert len(context.cws_results) == 3, (
|
||||
f"Expected 3 strategies, got {len(context.cws_results)}"
|
||||
)
|
||||
for _, conf in context.cws_results:
|
||||
assert conf > 0.0
|
||||
|
||||
|
||||
@then("the results should be sorted by confidence descending")
|
||||
def step_cws_sorted(context: Context) -> None:
|
||||
confs = [c for _, c in context.cws_results]
|
||||
assert confs == sorted(confs, reverse=True), (
|
||||
f"Results not sorted descending: {confs}"
|
||||
)
|
||||
|
||||
|
||||
@then('the "{name}" strategy should have a boosted confidence')
|
||||
def step_cws_boosted(context: Context, name: str) -> None:
|
||||
boosted = {s.name: c for s, c in context.cws_results_preferred}
|
||||
baseline = {s.name: c for s, c in context.cws_results_baseline}
|
||||
assert boosted[name] >= baseline[name], (
|
||||
f"Expected {name} boosted ({boosted[name]}) >= baseline ({baseline[name]})"
|
||||
)
|
||||
|
||||
|
||||
@then("the boosted confidence should not exceed 1.0")
|
||||
def step_cws_boosted_capped(context: Context) -> None:
|
||||
for _, conf in context.cws_results_preferred:
|
||||
assert conf <= 1.0, f"Confidence {conf} exceeds 1.0"
|
||||
|
||||
|
||||
@then("the zero-confidence strategy should not be in the results")
|
||||
def step_cws_no_zero(context: Context) -> None:
|
||||
names = [s.name for s, _ in context.cws_results_with_zero]
|
||||
assert "zero_conf" not in names, (
|
||||
f"Zero-confidence strategy should be excluded, got {names}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Then steps — ProportionalBudgetAllocator
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@then(
|
||||
"the first candidate should receive approximately {tokens:d} tokens from proportional allocator"
|
||||
)
|
||||
def step_pba_first(context: Context, tokens: int) -> None:
|
||||
actual = context.pba_allocations[0][2]
|
||||
assert abs(actual - tokens) <= 1, (
|
||||
f"First candidate got {actual}, expected ~{tokens}"
|
||||
)
|
||||
|
||||
|
||||
@then(
|
||||
"the second candidate should receive approximately {tokens:d} tokens from proportional allocator"
|
||||
)
|
||||
def step_pba_second(context: Context, tokens: int) -> None:
|
||||
actual = context.pba_allocations[1][2]
|
||||
assert abs(actual - tokens) <= 1, (
|
||||
f"Second candidate got {actual}, expected ~{tokens}"
|
||||
)
|
||||
|
||||
|
||||
@then("the total proportional allocation should equal exactly {budget:d}")
|
||||
def step_pba_total_exact(context: Context, budget: int) -> None:
|
||||
total = sum(a[2] for a in context.pba_allocations)
|
||||
assert total == budget, f"Total allocation {total} != {budget}"
|
||||
|
||||
|
||||
@then("only the high-confidence candidate should receive tokens")
|
||||
def step_pba_only_high(context: Context) -> None:
|
||||
assert len(context.pba_allocations) == 1, (
|
||||
f"Expected 1 allocation, got {len(context.pba_allocations)}"
|
||||
)
|
||||
|
||||
|
||||
@then("the single candidate should receive all {budget:d} tokens")
|
||||
def step_pba_single_all(context: Context, budget: int) -> None:
|
||||
assert len(context.pba_allocations) == 1
|
||||
assert context.pba_allocations[0][2] == budget
|
||||
|
||||
|
||||
@then("each proportional allocation should be {low:d} or {high:d}")
|
||||
def step_pba_each_bounded(context: Context, low: int, high: int) -> None:
|
||||
for _, _, tokens in context.pba_allocations:
|
||||
assert tokens in {low, high}, (
|
||||
f"Expected allocation to be {low} or {high}, got {tokens}"
|
||||
)
|
||||
|
||||
|
||||
@then("the highest-confidence candidate should receive the full budget")
|
||||
def step_pba_fallback(context: Context) -> None:
|
||||
# When all candidates are excluded, the allocator falls back to the
|
||||
# highest-confidence candidate with the full budget.
|
||||
total = sum(a[2] for a in context.pba_allocations)
|
||||
assert total == context.pba_budget, (
|
||||
f"Expected full budget {context.pba_budget}, got {total}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Then steps — CircuitBreaker
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@then('the circuit for "{name}" should be open')
|
||||
def step_cb_is_open(context: Context, name: str) -> None:
|
||||
assert context.circuit_breaker.is_open(name), (
|
||||
f"Expected circuit for '{name}' to be open"
|
||||
)
|
||||
|
||||
|
||||
@then('the circuit for "{name}" should be closed')
|
||||
def step_cb_is_closed(context: Context, name: str) -> None:
|
||||
assert not context.circuit_breaker.is_open(name), (
|
||||
f"Expected circuit for '{name}' to be closed"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Then steps — ParallelStrategyExecutor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@then("the executor should return fragments from the strategy")
|
||||
def step_executor_has_results(context: Context) -> None:
|
||||
assert len(context.executor_results) > 0, "Expected fragments from executor"
|
||||
|
||||
|
||||
@then("the tracking strategy should have been invoked")
|
||||
def step_tracking_invoked(context: Context) -> None:
|
||||
assert context.tracking_strategy.invoked, "Tracking strategy was not invoked"
|
||||
|
||||
|
||||
@then("the executor should return {count:d} fragments")
|
||||
def step_executor_count(context: Context, count: int) -> None:
|
||||
assert len(context.executor_results) == count, (
|
||||
f"Expected {count} fragments, got {len(context.executor_results)}"
|
||||
)
|
||||
|
||||
|
||||
@then('the circuit breaker should have recorded a failure for "{name}"')
|
||||
def step_cb_has_failure(context: Context, name: str) -> None:
|
||||
# After one failure, the circuit isn't necessarily open (depends on threshold)
|
||||
# but the failure count should be > 0
|
||||
assert context.executor_cb._failures.get(name, 0) > 0, (
|
||||
f"Expected failure recorded for '{name}'"
|
||||
)
|
||||
|
||||
|
||||
@then("the executor should return fragments from all parallel strategies")
|
||||
def step_executor_parallel_results(context: Context) -> None:
|
||||
# Two strategies each returning up to their budget from the same 2 fragments
|
||||
assert len(context.executor_results) > 0, (
|
||||
"Expected fragments from parallel execution"
|
||||
)
|
||||
|
||||
|
||||
@then("both parallel strategies should have been invoked")
|
||||
def step_parallel_both_invoked(context: Context) -> None:
|
||||
for s in context.parallel_strategies:
|
||||
assert s.invoked, f"Strategy {s.name} was not invoked"
|
||||
|
||||
|
||||
@then("the executor should return fragments only from the successful parallel strategy")
|
||||
def step_executor_parallel_mixed_results(context: Context) -> None:
|
||||
assert len(context.executor_results) > 0, (
|
||||
"Expected fragments from the successful strategy"
|
||||
)
|
||||
assert context.parallel_ok_strategy.invoked, (
|
||||
"Good strategy should have been invoked"
|
||||
)
|
||||
|
||||
|
||||
@then('the parallel circuit breaker should record a failure for "{name}"')
|
||||
def step_parallel_cb_failure(context: Context, name: str) -> None:
|
||||
assert context.parallel_cb._failures.get(name, 0) > 0, (
|
||||
f"Expected failure recorded for '{name}'"
|
||||
)
|
||||
|
||||
|
||||
@then("the circuit_breaker property should return the custom breaker")
|
||||
def step_executor_cb_property(context: Context) -> None:
|
||||
assert context.custom_executor.circuit_breaker is context.custom_cb, (
|
||||
"circuit_breaker property did not return the custom breaker"
|
||||
)
|
||||
|
||||
|
||||
@then("the proportional allocator should return an empty list")
|
||||
def step_pba_empty_result(context: Context) -> None:
|
||||
assert len(context.pba_allocations) == 0, (
|
||||
f"Expected empty allocations, got {len(context.pba_allocations)}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Then steps — ContextAssemblyPipeline
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@then("the orchestrator payload should contain {count:d} fragments")
|
||||
def step_orch_payload_count(context: Context, count: int) -> None:
|
||||
assert len(context.orch_payload.fragments) == count, (
|
||||
f"Expected {count} fragments, got {len(context.orch_payload.fragments)}"
|
||||
)
|
||||
|
||||
|
||||
@then("the orchestrator should have timing metadata")
|
||||
def step_orch_has_timings(context: Context) -> None:
|
||||
timings = context.orch_pipeline.last_timings
|
||||
assert timings is not None, "Expected timing metadata to be populated"
|
||||
context.orch_timings = timings
|
||||
|
||||
|
||||
@then("all timing values should be non-negative")
|
||||
def step_orch_timings_nonneg(context: Context) -> None:
|
||||
timings = context.orch_pipeline.last_timings
|
||||
assert timings is not None
|
||||
for name in type(timings).model_fields:
|
||||
val = getattr(timings, name)
|
||||
assert val >= 0.0, f"Timing {name} = {val} is negative"
|
||||
|
||||
|
||||
@then('the orchestrator payload strategies used should include "{strategy}"')
|
||||
def step_orch_strategies(context: Context, strategy: str) -> None:
|
||||
assert strategy in context.orch_payload.strategies_used, (
|
||||
f"Expected {strategy!r} in {context.orch_payload.strategies_used}"
|
||||
)
|
||||
|
||||
|
||||
@then("the orchestrator payload should be within budget")
|
||||
def step_orch_within_budget(context: Context) -> None:
|
||||
assert context.orch_payload.is_within_budget
|
||||
|
||||
|
||||
@then('an orchestrator ValueError should be raised mentioning "{keyword}"')
|
||||
def step_orch_error(context: Context, keyword: str) -> None:
|
||||
assert context.orch_error is not None, "Expected ValueError but none raised"
|
||||
assert keyword.lower() in str(context.orch_error).lower(), (
|
||||
f"Expected '{keyword}' in error: {context.orch_error}"
|
||||
)
|
||||
|
||||
|
||||
@then("the custom orchestrator selector should have been called")
|
||||
def step_orch_selector_called(context: Context) -> None:
|
||||
assert context.orch_custom_selector.called, (
|
||||
"Custom orchestrator selector was not invoked"
|
||||
)
|
||||
|
||||
|
||||
@then('the first orchestrator fragment content should be "{content}"')
|
||||
def step_orch_first_content(context: Context, content: str) -> None:
|
||||
assert context.orch_payload.fragments[0].content == content, (
|
||||
f"Expected first fragment '{content}', "
|
||||
f"got '{context.orch_payload.fragments[0].content}'"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Then steps — StageTimings
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@then("the timings total_ms should be {total:g}")
|
||||
def step_timings_total(context: Context, total: float) -> None:
|
||||
assert context.test_timings.total_ms == total
|
||||
|
||||
|
||||
@then("the timings should have all 10 named fields")
|
||||
def step_timings_fields(context: Context) -> None:
|
||||
field_names = set(StageTimings.model_fields.keys())
|
||||
assert len(field_names) == 10, f"Expected 10 fields, got {len(field_names)}"
|
||||
expected_names = {
|
||||
"strategy_selection_ms",
|
||||
"budget_allocation_ms",
|
||||
"strategy_execution_ms",
|
||||
"deduplication_ms",
|
||||
"depth_resolution_ms",
|
||||
"scoring_ms",
|
||||
"packing_ms",
|
||||
"ordering_ms",
|
||||
"preamble_generation_ms",
|
||||
"total_ms",
|
||||
}
|
||||
assert field_names == expected_names, (
|
||||
f"Field mismatch: {field_names.symmetric_difference(expected_names)}"
|
||||
)
|
||||
@@ -0,0 +1,81 @@
|
||||
*** Settings ***
|
||||
Documentation Integration smoke tests for ACMS Pipeline Orchestrator and Phase 1 components
|
||||
Resource ${CURDIR}/common.resource
|
||||
Suite Setup Setup Test Environment
|
||||
Suite Teardown Cleanup Test Environment
|
||||
|
||||
*** Variables ***
|
||||
${HELPER} ${CURDIR}/helper_acms_pipeline_orchestrator.py
|
||||
|
||||
*** Test Cases ***
|
||||
ConfidenceWeightedSelector Selects Strategies
|
||||
[Documentation] ConfidenceWeightedSelector selects and ranks strategies by confidence
|
||||
${result}= Run Process ${PYTHON} ${HELPER} cws-select cwd=${WORKSPACE}
|
||||
Log ${result.stdout}
|
||||
Log ${result.stderr}
|
||||
Should Be Equal As Integers ${result.rc} 0
|
||||
Should Contain ${result.stdout} cws-select-ok
|
||||
|
||||
ConfidenceWeightedSelector Boosts Preferred
|
||||
[Documentation] ConfidenceWeightedSelector boosts preferred strategies
|
||||
${result}= Run Process ${PYTHON} ${HELPER} cws-preferred cwd=${WORKSPACE}
|
||||
Log ${result.stdout}
|
||||
Log ${result.stderr}
|
||||
Should Be Equal As Integers ${result.rc} 0
|
||||
Should Contain ${result.stdout} cws-preferred-ok
|
||||
|
||||
ProportionalBudgetAllocator Distributes Budget
|
||||
[Documentation] ProportionalBudgetAllocator distributes tokens proportionally
|
||||
${result}= Run Process ${PYTHON} ${HELPER} pba-proportional cwd=${WORKSPACE}
|
||||
Log ${result.stdout}
|
||||
Log ${result.stderr}
|
||||
Should Be Equal As Integers ${result.rc} 0
|
||||
Should Contain ${result.stdout} pba-proportional-ok
|
||||
|
||||
ProportionalBudgetAllocator MinBudget Enforcement
|
||||
[Documentation] ProportionalBudgetAllocator enforces min_useful_budget
|
||||
${result}= Run Process ${PYTHON} ${HELPER} pba-min-budget cwd=${WORKSPACE}
|
||||
Log ${result.stdout}
|
||||
Log ${result.stderr}
|
||||
Should Be Equal As Integers ${result.rc} 0
|
||||
Should Contain ${result.stdout} pba-min-budget-ok
|
||||
|
||||
CircuitBreaker Opens After Failures
|
||||
[Documentation] CircuitBreaker opens circuit after consecutive failures
|
||||
${result}= Run Process ${PYTHON} ${HELPER} cb-open cwd=${WORKSPACE}
|
||||
Log ${result.stdout}
|
||||
Log ${result.stderr}
|
||||
Should Be Equal As Integers ${result.rc} 0
|
||||
Should Contain ${result.stdout} cb-open-ok
|
||||
|
||||
CircuitBreaker Resets On Success
|
||||
[Documentation] CircuitBreaker resets on success
|
||||
${result}= Run Process ${PYTHON} ${HELPER} cb-reset cwd=${WORKSPACE}
|
||||
Log ${result.stdout}
|
||||
Log ${result.stderr}
|
||||
Should Be Equal As Integers ${result.rc} 0
|
||||
Should Contain ${result.stdout} cb-reset-ok
|
||||
|
||||
ParallelStrategyExecutor Executes Strategy
|
||||
[Documentation] ParallelStrategyExecutor invokes strategy and returns fragments
|
||||
${result}= Run Process ${PYTHON} ${HELPER} executor-run cwd=${WORKSPACE}
|
||||
Log ${result.stdout}
|
||||
Log ${result.stderr}
|
||||
Should Be Equal As Integers ${result.rc} 0
|
||||
Should Contain ${result.stdout} executor-run-ok
|
||||
|
||||
ContextAssemblyPipeline Assemble With Timings
|
||||
[Documentation] Full pipeline assembly with timing metadata
|
||||
${result}= Run Process ${PYTHON} ${HELPER} pipeline-assemble cwd=${WORKSPACE}
|
||||
Log ${result.stdout}
|
||||
Log ${result.stderr}
|
||||
Should Be Equal As Integers ${result.rc} 0
|
||||
Should Contain ${result.stdout} pipeline-assemble-ok
|
||||
|
||||
ContextAssemblyPipeline Custom Components
|
||||
[Documentation] Pipeline accepts custom Phase 1 components
|
||||
${result}= Run Process ${PYTHON} ${HELPER} pipeline-custom cwd=${WORKSPACE}
|
||||
Log ${result.stdout}
|
||||
Log ${result.stderr}
|
||||
Should Be Equal As Integers ${result.rc} 0
|
||||
Should Contain ${result.stdout} pipeline-custom-ok
|
||||
@@ -0,0 +1,299 @@
|
||||
"""Robot Framework helper for ACMS Pipeline Orchestrator integration tests.
|
||||
|
||||
Provides a CLI-style interface for Robot to invoke Phase 1 pipeline
|
||||
component operations and verify the results.
|
||||
|
||||
Usage:
|
||||
python robot/helper_acms_pipeline_orchestrator.py cws-select
|
||||
python robot/helper_acms_pipeline_orchestrator.py cws-preferred
|
||||
python robot/helper_acms_pipeline_orchestrator.py pba-proportional
|
||||
python robot/helper_acms_pipeline_orchestrator.py pba-min-budget
|
||||
python robot/helper_acms_pipeline_orchestrator.py cb-open
|
||||
python robot/helper_acms_pipeline_orchestrator.py cb-reset
|
||||
python robot/helper_acms_pipeline_orchestrator.py executor-run
|
||||
python robot/helper_acms_pipeline_orchestrator.py pipeline-assemble
|
||||
python robot/helper_acms_pipeline_orchestrator.py pipeline-custom
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from collections.abc import Callable, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
_SRC = str(Path(__file__).resolve().parents[1] / "src")
|
||||
if _SRC not in sys.path:
|
||||
sys.path.insert(0, _SRC)
|
||||
|
||||
from cleveragents.application.services.acms_pipeline import ( # noqa: E402
|
||||
CircuitBreaker,
|
||||
ConfidenceWeightedSelector,
|
||||
ContextAssemblyPipeline,
|
||||
ParallelStrategyExecutor,
|
||||
ProportionalBudgetAllocator,
|
||||
)
|
||||
from cleveragents.application.services.acms_service import ( # noqa: E402
|
||||
DefaultStrategySelector,
|
||||
RecencyStrategy,
|
||||
RelevanceStrategy,
|
||||
StrategyCapabilities,
|
||||
TieredStrategy,
|
||||
)
|
||||
from cleveragents.domain.models.core.context_fragment import ( # noqa: E402
|
||||
ContextBudget,
|
||||
ContextFragment,
|
||||
FragmentProvenance,
|
||||
)
|
||||
|
||||
_DEFAULT_PROV = FragmentProvenance(resource_uri="test://robot-orch")
|
||||
|
||||
|
||||
def _make_frag(**kwargs: Any) -> ContextFragment:
|
||||
kwargs.setdefault("uko_node", "test://robot-orch")
|
||||
kwargs.setdefault("token_count", 0)
|
||||
kwargs.setdefault("provenance", _DEFAULT_PROV)
|
||||
return ContextFragment(**kwargs)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Commands
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _cmd_cws_select() -> int:
|
||||
"""ConfidenceWeightedSelector selects strategies."""
|
||||
strategies = [RelevanceStrategy(), RecencyStrategy(), TieredStrategy()]
|
||||
selector = ConfidenceWeightedSelector()
|
||||
results = selector.select(strategies, {})
|
||||
assert len(results) == 3, f"Expected 3, got {len(results)}"
|
||||
confs = [c for _, c in results]
|
||||
assert confs == sorted(confs, reverse=True), "Not sorted"
|
||||
print("cws-select-ok")
|
||||
return 0
|
||||
|
||||
|
||||
def _cmd_cws_preferred() -> int:
|
||||
"""ConfidenceWeightedSelector boosts preferred strategies."""
|
||||
strategies = [RelevanceStrategy(), RecencyStrategy(), TieredStrategy()]
|
||||
selector = ConfidenceWeightedSelector(preference_boost=1.5)
|
||||
baseline = dict(selector.select(strategies, {}))
|
||||
boosted = dict(selector.select(strategies, {"preferred_strategies": ["recency"]}))
|
||||
# recency should be boosted
|
||||
recency_baseline = baseline.get(
|
||||
next(s for s in strategies if s.name == "recency"), 0.0
|
||||
)
|
||||
recency_boosted = boosted.get(
|
||||
next(s for s in strategies if s.name == "recency"), 0.0
|
||||
)
|
||||
assert recency_boosted >= recency_baseline
|
||||
assert recency_boosted <= 1.0
|
||||
print("cws-preferred-ok")
|
||||
return 0
|
||||
|
||||
|
||||
def _cmd_pba_proportional() -> int:
|
||||
"""ProportionalBudgetAllocator distributes proportionally."""
|
||||
s1 = RelevanceStrategy()
|
||||
s2 = RecencyStrategy()
|
||||
allocator = ProportionalBudgetAllocator()
|
||||
allocs = allocator.allocate([(s1, 0.8), (s2, 0.2)], 1000)
|
||||
total = sum(a[2] for a in allocs)
|
||||
assert total == 1000, f"Total {total} != 1000"
|
||||
assert abs(allocs[0][2] - 800) <= 1
|
||||
assert abs(allocs[1][2] - 200) <= 1
|
||||
print("pba-proportional-ok")
|
||||
return 0
|
||||
|
||||
|
||||
def _cmd_pba_min_budget() -> int:
|
||||
"""ProportionalBudgetAllocator enforces min_useful_budget."""
|
||||
s1 = RelevanceStrategy()
|
||||
s2 = RecencyStrategy()
|
||||
allocator = ProportionalBudgetAllocator(min_useful_budget=200)
|
||||
allocs = allocator.allocate([(s1, 0.99), (s2, 0.01)], 300)
|
||||
assert len(allocs) == 1, f"Expected 1, got {len(allocs)}"
|
||||
assert allocs[0][2] == 300
|
||||
print("pba-min-budget-ok")
|
||||
return 0
|
||||
|
||||
|
||||
def _cmd_cb_open() -> int:
|
||||
"""CircuitBreaker opens after threshold failures."""
|
||||
cb = CircuitBreaker(failure_threshold=3)
|
||||
for _ in range(3):
|
||||
cb.record_failure("flaky")
|
||||
assert cb.is_open("flaky"), "Circuit should be open"
|
||||
print("cb-open-ok")
|
||||
return 0
|
||||
|
||||
|
||||
def _cmd_cb_reset() -> int:
|
||||
"""CircuitBreaker resets on success."""
|
||||
cb = CircuitBreaker(failure_threshold=3)
|
||||
cb.record_failure("recovering")
|
||||
cb.record_failure("recovering")
|
||||
cb.record_success("recovering")
|
||||
assert not cb.is_open("recovering"), "Circuit should be closed"
|
||||
print("cb-reset-ok")
|
||||
return 0
|
||||
|
||||
|
||||
class _TrackingRobotStrategy:
|
||||
"""Test strategy that tracks invocation."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.invoked = False
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "tracking_robot"
|
||||
|
||||
@property
|
||||
def capabilities(self) -> StrategyCapabilities:
|
||||
return StrategyCapabilities()
|
||||
|
||||
def can_handle(self, request: dict[str, Any]) -> float:
|
||||
return 0.7
|
||||
|
||||
def assemble(
|
||||
self,
|
||||
fragments: Sequence[ContextFragment],
|
||||
budget: ContextBudget,
|
||||
) -> Sequence[ContextFragment]:
|
||||
self.invoked = True
|
||||
result: list[ContextFragment] = []
|
||||
total = 0
|
||||
for frag in fragments:
|
||||
if total + frag.token_count <= budget.available_tokens:
|
||||
result.append(frag)
|
||||
total += frag.token_count
|
||||
return result
|
||||
|
||||
def explain(self) -> str:
|
||||
return "Tracking."
|
||||
|
||||
|
||||
def _cmd_executor_run() -> int:
|
||||
"""ParallelStrategyExecutor invokes strategy."""
|
||||
strategy = _TrackingRobotStrategy()
|
||||
frags = [
|
||||
_make_frag(
|
||||
uko_node="project://t/a.py",
|
||||
content="alpha",
|
||||
token_count=50,
|
||||
relevance_score=0.9,
|
||||
),
|
||||
]
|
||||
executor = ParallelStrategyExecutor()
|
||||
budget = ContextBudget(max_tokens=200, reserved_tokens=0)
|
||||
results = executor.execute([(strategy, 0.7, 200)], frags, budget)
|
||||
assert len(results) > 0, "Expected fragments"
|
||||
assert strategy.invoked, "Strategy not invoked"
|
||||
print("executor-run-ok")
|
||||
return 0
|
||||
|
||||
|
||||
def _cmd_pipeline_assemble() -> int:
|
||||
"""ContextAssemblyPipeline assembly with timing."""
|
||||
pipeline = ContextAssemblyPipeline()
|
||||
frags = [
|
||||
_make_frag(
|
||||
uko_node="project://t/a.py",
|
||||
content="alpha",
|
||||
token_count=100,
|
||||
relevance_score=0.9,
|
||||
),
|
||||
_make_frag(
|
||||
uko_node="project://t/b.py",
|
||||
content="beta",
|
||||
token_count=100,
|
||||
relevance_score=0.3,
|
||||
),
|
||||
]
|
||||
budget = ContextBudget(max_tokens=250, reserved_tokens=0)
|
||||
payload = pipeline.assemble(
|
||||
plan_id="01JQTESTPN00000000000000AA",
|
||||
fragments=frags,
|
||||
budget=budget,
|
||||
strategy="relevance",
|
||||
)
|
||||
assert len(payload.fragments) == 2
|
||||
timings = pipeline.last_timings
|
||||
assert timings is not None
|
||||
assert timings.total_ms >= 0.0
|
||||
for name in type(timings).model_fields:
|
||||
assert getattr(timings, name) >= 0.0
|
||||
print("pipeline-assemble-ok")
|
||||
return 0
|
||||
|
||||
|
||||
class _TrackingSelector:
|
||||
"""Tracks invocation."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.called = False
|
||||
|
||||
def select(
|
||||
self,
|
||||
strategies: Sequence[Any],
|
||||
request: dict[str, Any],
|
||||
) -> list[tuple[Any, float]]:
|
||||
self.called = True
|
||||
return DefaultStrategySelector().select(strategies, request)
|
||||
|
||||
|
||||
def _cmd_pipeline_custom() -> int:
|
||||
"""ContextAssemblyPipeline with custom components."""
|
||||
selector = _TrackingSelector()
|
||||
pipeline = ContextAssemblyPipeline(strategy_selector=selector)
|
||||
frags = [
|
||||
_make_frag(
|
||||
uko_node="project://t/a.py",
|
||||
content="alpha",
|
||||
token_count=100,
|
||||
relevance_score=0.9,
|
||||
),
|
||||
]
|
||||
budget = ContextBudget(max_tokens=500, reserved_tokens=0)
|
||||
pipeline.assemble(
|
||||
plan_id="01JQTESTPN00000000000000AA",
|
||||
fragments=frags,
|
||||
budget=budget,
|
||||
strategy="relevance",
|
||||
)
|
||||
assert selector.called, "Custom selector not invoked"
|
||||
print("pipeline-custom-ok")
|
||||
return 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI dispatcher
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_COMMANDS: dict[str, Callable[[], int]] = {
|
||||
"cws-select": _cmd_cws_select,
|
||||
"cws-preferred": _cmd_cws_preferred,
|
||||
"pba-proportional": _cmd_pba_proportional,
|
||||
"pba-min-budget": _cmd_pba_min_budget,
|
||||
"cb-open": _cmd_cb_open,
|
||||
"cb-reset": _cmd_cb_reset,
|
||||
"executor-run": _cmd_executor_run,
|
||||
"pipeline-assemble": _cmd_pipeline_assemble,
|
||||
"pipeline-custom": _cmd_pipeline_custom,
|
||||
}
|
||||
|
||||
|
||||
def main() -> int:
|
||||
if len(sys.argv) < 2 or sys.argv[1] not in _COMMANDS:
|
||||
print(f"Usage: {sys.argv[0]} <{'|'.join(_COMMANDS)}>", file=sys.stderr)
|
||||
return 1
|
||||
try:
|
||||
return _COMMANDS[sys.argv[1]]()
|
||||
except Exception as exc:
|
||||
print(f"FAIL: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,679 @@
|
||||
"""ACMS Context Assembly Pipeline orchestrator and Phase 1 components.
|
||||
|
||||
Implements the production-quality versions of the Phase 1 pipeline components
|
||||
defined in ``docs/specification.md`` §42630-42636:
|
||||
|
||||
- **ConfidenceWeightedSelector** — Selects strategies by confidence score
|
||||
weighted by backend availability. Respects ``preferred_strategies`` hints
|
||||
from the ``ContextRequest``.
|
||||
|
||||
- **ProportionalBudgetAllocator** — Distributes token budget proportionally
|
||||
across selected strategies with ``min_useful_budget`` enforcement.
|
||||
|
||||
- **ParallelStrategyExecutor** — Executes strategies concurrently via
|
||||
``ThreadPoolExecutor`` with per-strategy timeouts and a circuit-breaker
|
||||
pattern (3 consecutive failures → temporarily disabled).
|
||||
|
||||
- **ContextAssemblyPipeline** — Full 10-stage pipeline mediator that wires
|
||||
all components and adds structured logging with per-stage timing.
|
||||
|
||||
All components conform to the Protocol interfaces defined in
|
||||
``acms_service.py`` and can be used as drop-in replacements for the
|
||||
``Default*`` pass-through stubs.
|
||||
|
||||
Based on ``docs/specification.md`` ~line 42615.
|
||||
|
||||
ISSUES CLOSED: #539
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import structlog
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from cleveragents.config.settings import Settings
|
||||
from cleveragents.infrastructure.database.unit_of_work import UnitOfWork
|
||||
|
||||
from cleveragents.application.services.acms_service import (
|
||||
ACMSPipeline,
|
||||
BudgetAllocator,
|
||||
BudgetPacker,
|
||||
ContextStrategy,
|
||||
DefaultBudgetPacker,
|
||||
DefaultDeduplicator,
|
||||
DefaultDepthResolver,
|
||||
DefaultOrderer,
|
||||
DefaultPreambleGenerator,
|
||||
DefaultScorer,
|
||||
DefaultSkeletonCompressor,
|
||||
DetailDepthResolver,
|
||||
FragmentDeduplicator,
|
||||
FragmentOrderer,
|
||||
FragmentScorer,
|
||||
PreambleGenerator,
|
||||
SkeletonCompressor,
|
||||
StrategyExecutor,
|
||||
StrategySelector,
|
||||
)
|
||||
from cleveragents.domain.models.core.context_fragment import (
|
||||
ContextBudget,
|
||||
ContextFragment,
|
||||
ContextPayload,
|
||||
build_provenance_map,
|
||||
compute_context_hash,
|
||||
)
|
||||
|
||||
logger = structlog.get_logger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Circuit Breaker (per-strategy failure tracking)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class CircuitBreaker:
|
||||
"""Per-strategy circuit breaker for the ``ParallelStrategyExecutor``.
|
||||
|
||||
Tracks consecutive failures per strategy name. After
|
||||
``failure_threshold`` consecutive failures the circuit *opens* and
|
||||
the strategy is temporarily disabled.
|
||||
|
||||
Thread-safe: internal state is only mutated from the executor thread
|
||||
that owns the futures, never concurrently.
|
||||
|
||||
Based on ``docs/specification.md`` ~line 42936 (ParallelStrategyExecutor
|
||||
uses Circuit Breaker pattern).
|
||||
"""
|
||||
|
||||
def __init__(self, failure_threshold: int = 3) -> None:
|
||||
self.failure_threshold = failure_threshold
|
||||
self._failures: dict[str, int] = {}
|
||||
self._disabled: set[str] = set()
|
||||
|
||||
def record_success(self, name: str) -> None:
|
||||
"""Record a successful execution, resetting the failure counter."""
|
||||
self._failures[name] = 0
|
||||
self._disabled.discard(name)
|
||||
|
||||
def record_failure(self, name: str) -> None:
|
||||
"""Record a failure; open the circuit after threshold breaches."""
|
||||
count = self._failures.get(name, 0) + 1
|
||||
self._failures[name] = count
|
||||
if count >= self.failure_threshold:
|
||||
self._disabled.add(name)
|
||||
|
||||
def is_open(self, name: str) -> bool:
|
||||
"""Return ``True`` if the strategy's circuit is open (disabled)."""
|
||||
return name in self._disabled
|
||||
|
||||
def reset(self, name: str) -> None:
|
||||
"""Reset a strategy's failure counter and re-enable it."""
|
||||
self._failures.pop(name, None)
|
||||
self._disabled.discard(name)
|
||||
|
||||
def reset_all(self) -> None:
|
||||
"""Reset all strategies."""
|
||||
self._failures.clear()
|
||||
self._disabled.clear()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 1 — ConfidenceWeightedSelector
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ConfidenceWeightedSelector:
|
||||
"""Select strategies by confidence score weighted by backend availability.
|
||||
|
||||
Evaluates all strategies via ``can_handle``, filters to confidence > 0,
|
||||
and sorts by effective confidence descending. When the request dict
|
||||
contains a ``preferred_strategies`` key (list of names), those
|
||||
strategies receive a 1.5x confidence boost.
|
||||
|
||||
Emits structured log events for strategy selection diagnostics.
|
||||
|
||||
Based on ``docs/specification.md`` ~line 42918
|
||||
(``ConfidenceWeightedSelector``).
|
||||
"""
|
||||
|
||||
def __init__(self, *, preference_boost: float = 1.5) -> None:
|
||||
self._preference_boost = preference_boost
|
||||
self._logger = logger.bind(component="ConfidenceWeightedSelector")
|
||||
|
||||
def select(
|
||||
self,
|
||||
strategies: Sequence[ContextStrategy],
|
||||
request: dict[str, Any],
|
||||
) -> list[tuple[ContextStrategy, float]]:
|
||||
"""Return (strategy, confidence) pairs sorted by priority."""
|
||||
preferred: list[str] = request.get("preferred_strategies", [])
|
||||
scored: list[tuple[ContextStrategy, float]] = []
|
||||
|
||||
for strategy in strategies:
|
||||
raw_confidence = strategy.can_handle(request)
|
||||
if raw_confidence <= 0.0:
|
||||
continue
|
||||
effective = raw_confidence
|
||||
if preferred and strategy.name in preferred:
|
||||
effective = min(raw_confidence * self._preference_boost, 1.0)
|
||||
scored.append((strategy, effective))
|
||||
|
||||
scored.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
self._logger.debug(
|
||||
"Strategy selection complete",
|
||||
candidates=len(scored),
|
||||
preferred=preferred,
|
||||
)
|
||||
return scored
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 1 — ProportionalBudgetAllocator
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
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.
|
||||
|
||||
Uses the largest-remainder method to ensure the full budget is consumed
|
||||
with no token loss.
|
||||
|
||||
Based on ``docs/specification.md`` ~line 42918
|
||||
(``ProportionalBudgetAllocator``).
|
||||
"""
|
||||
|
||||
def __init__(self, *, min_useful_budget: int = 64) -> None:
|
||||
self._min_useful_budget = min_useful_budget
|
||||
self._logger = logger.bind(component="ProportionalBudgetAllocator")
|
||||
|
||||
def allocate(
|
||||
self,
|
||||
candidates: list[tuple[ContextStrategy, float]],
|
||||
total_budget: int,
|
||||
) -> list[tuple[ContextStrategy, float, int]]:
|
||||
"""Return (strategy, confidence, allocated_tokens) triples."""
|
||||
if not candidates:
|
||||
return []
|
||||
|
||||
# Filter candidates below min_useful_budget.
|
||||
# Single candidate always gets the full budget.
|
||||
working = list(candidates)
|
||||
if len(working) > 1:
|
||||
working = self._filter_by_min_budget(working, total_budget)
|
||||
|
||||
if not working:
|
||||
# All candidates excluded — fall back to the highest-confidence one
|
||||
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
|
||||
return self._equal_split(working, total_budget)
|
||||
|
||||
return self._proportional_split(working, total_budget, total_confidence)
|
||||
|
||||
def _filter_by_min_budget(
|
||||
self,
|
||||
candidates: list[tuple[ContextStrategy, float]],
|
||||
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:
|
||||
return candidates # Can't compute proportions
|
||||
kept: list[tuple[ContextStrategy, float]] = []
|
||||
for strategy, confidence in candidates:
|
||||
share = int(total_budget * confidence / total_conf)
|
||||
if share >= self._min_useful_budget:
|
||||
kept.append((strategy, confidence))
|
||||
if not kept:
|
||||
return candidates # Don't exclude everyone
|
||||
return kept
|
||||
|
||||
def _equal_split(
|
||||
self,
|
||||
candidates: list[tuple[ContextStrategy, float]],
|
||||
total_budget: int,
|
||||
) -> list[tuple[ContextStrategy, float, int]]:
|
||||
"""Split budget equally with largest-remainder distribution."""
|
||||
n = len(candidates)
|
||||
share = total_budget // n
|
||||
remainder = total_budget - share * n
|
||||
return [
|
||||
(s, c, share + (1 if i < remainder else 0))
|
||||
for i, (s, c) in enumerate(candidates)
|
||||
]
|
||||
|
||||
def _proportional_split(
|
||||
self,
|
||||
candidates: list[tuple[ContextStrategy, float]],
|
||||
total_budget: int,
|
||||
total_confidence: float,
|
||||
) -> list[tuple[ContextStrategy, float, int]]:
|
||||
"""Proportional allocation with largest-remainder rounding."""
|
||||
raw = [(total_budget * c / total_confidence) for _, c in candidates]
|
||||
floors = [int(r) for r in raw]
|
||||
remainder = total_budget - sum(floors)
|
||||
fractions = [
|
||||
(r - f, i) for i, (r, f) in enumerate(zip(raw, floors, strict=True))
|
||||
]
|
||||
fractions.sort(reverse=True)
|
||||
for _, i in fractions[:remainder]:
|
||||
floors[i] += 1
|
||||
|
||||
self._logger.debug(
|
||||
"Budget allocation complete",
|
||||
candidates=len(candidates),
|
||||
total_budget=total_budget,
|
||||
allocations=floors,
|
||||
)
|
||||
return [(s, c, floors[i]) for i, (s, c) in enumerate(candidates)]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 1 — ParallelStrategyExecutor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ParallelStrategyExecutor:
|
||||
"""Execute strategies concurrently with timeouts and circuit breaking.
|
||||
|
||||
Uses ``ThreadPoolExecutor`` to run strategies in parallel. Each
|
||||
strategy receives a per-strategy timeout (default 30s). After
|
||||
``circuit_breaker.failure_threshold`` consecutive failures, a strategy
|
||||
is temporarily disabled for the current pipeline instance.
|
||||
|
||||
Failed strategies log warnings and are excluded from results — they do
|
||||
not cause the pipeline to fail.
|
||||
|
||||
Based on ``docs/specification.md`` ~line 42918
|
||||
(``ParallelStrategyExecutor``).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
timeout_seconds: float = 30.0,
|
||||
max_workers: int = 4,
|
||||
circuit_breaker: CircuitBreaker | None = None,
|
||||
) -> None:
|
||||
self._timeout_seconds = timeout_seconds
|
||||
self._max_workers = max_workers
|
||||
self._circuit_breaker = circuit_breaker or CircuitBreaker()
|
||||
self._logger = logger.bind(component="ParallelStrategyExecutor")
|
||||
|
||||
@property
|
||||
def circuit_breaker(self) -> CircuitBreaker:
|
||||
"""Access the circuit breaker for inspection/reset."""
|
||||
return self._circuit_breaker
|
||||
|
||||
def execute(
|
||||
self,
|
||||
allocations: list[tuple[ContextStrategy, float, int]],
|
||||
fragments: Sequence[ContextFragment],
|
||||
budget: ContextBudget,
|
||||
) -> Sequence[ContextFragment]:
|
||||
"""Execute strategies and return collected fragments."""
|
||||
if not allocations:
|
||||
return []
|
||||
|
||||
# Filter out circuit-broken and zero-budget strategies
|
||||
active: list[tuple[ContextStrategy, float, int]] = []
|
||||
for strategy, confidence, allocated_tokens in allocations:
|
||||
if self._circuit_breaker.is_open(strategy.name):
|
||||
self._logger.warning(
|
||||
"Strategy circuit-broken, skipping",
|
||||
strategy=strategy.name,
|
||||
)
|
||||
continue
|
||||
if allocated_tokens <= 0:
|
||||
continue
|
||||
active.append((strategy, confidence, allocated_tokens))
|
||||
|
||||
if not active:
|
||||
return []
|
||||
|
||||
# For single strategy, avoid threading overhead
|
||||
if len(active) == 1:
|
||||
return self._execute_single(active[0], fragments)
|
||||
|
||||
return self._execute_parallel(active, fragments)
|
||||
|
||||
def _execute_single(
|
||||
self,
|
||||
allocation: tuple[ContextStrategy, float, int],
|
||||
fragments: Sequence[ContextFragment],
|
||||
) -> list[ContextFragment]:
|
||||
"""Execute a single strategy synchronously."""
|
||||
strategy, _confidence, allocated_tokens = allocation
|
||||
scoped_budget = ContextBudget(
|
||||
max_tokens=allocated_tokens,
|
||||
reserved_tokens=0,
|
||||
)
|
||||
start = time.monotonic()
|
||||
try:
|
||||
result = list(strategy.assemble(fragments, scoped_budget))
|
||||
elapsed = time.monotonic() - start
|
||||
self._circuit_breaker.record_success(strategy.name)
|
||||
self._logger.info(
|
||||
"Strategy executed",
|
||||
strategy=strategy.name,
|
||||
fragments_returned=len(result),
|
||||
elapsed_ms=round(elapsed * 1000, 1),
|
||||
)
|
||||
return result
|
||||
except Exception:
|
||||
elapsed = time.monotonic() - start
|
||||
self._circuit_breaker.record_failure(strategy.name)
|
||||
self._logger.warning(
|
||||
"Strategy failed",
|
||||
strategy=strategy.name,
|
||||
elapsed_ms=round(elapsed * 1000, 1),
|
||||
exc_info=True,
|
||||
)
|
||||
return []
|
||||
|
||||
def _execute_parallel(
|
||||
self,
|
||||
active: list[tuple[ContextStrategy, float, int]],
|
||||
fragments: Sequence[ContextFragment],
|
||||
) -> list[ContextFragment]:
|
||||
"""Execute multiple strategies in parallel via ThreadPoolExecutor."""
|
||||
collected: list[ContextFragment] = []
|
||||
|
||||
with ThreadPoolExecutor(
|
||||
max_workers=min(self._max_workers, len(active))
|
||||
) as executor:
|
||||
future_to_name: dict[Any, str] = {}
|
||||
for strategy, _confidence, allocated_tokens in active:
|
||||
scoped_budget = ContextBudget(
|
||||
max_tokens=allocated_tokens,
|
||||
reserved_tokens=0,
|
||||
)
|
||||
future = executor.submit(
|
||||
self._run_strategy, strategy, fragments, scoped_budget
|
||||
)
|
||||
future_to_name[future] = strategy.name
|
||||
|
||||
for future in as_completed(future_to_name, timeout=self._timeout_seconds):
|
||||
name = future_to_name[future]
|
||||
try:
|
||||
result = future.result(timeout=0)
|
||||
self._circuit_breaker.record_success(name)
|
||||
collected.extend(result)
|
||||
self._logger.info(
|
||||
"Strategy executed (parallel)",
|
||||
strategy=name,
|
||||
fragments_returned=len(result),
|
||||
)
|
||||
except Exception:
|
||||
self._circuit_breaker.record_failure(name)
|
||||
self._logger.warning(
|
||||
"Strategy failed (parallel)",
|
||||
strategy=name,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
return collected
|
||||
|
||||
@staticmethod
|
||||
def _run_strategy(
|
||||
strategy: ContextStrategy,
|
||||
fragments: Sequence[ContextFragment],
|
||||
budget: ContextBudget,
|
||||
) -> list[ContextFragment]:
|
||||
"""Invoke a single strategy — called from worker thread."""
|
||||
return list(strategy.assemble(fragments, budget))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Timing helper
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class StageTimings(BaseModel, frozen=True):
|
||||
"""Timing data for each pipeline stage, in milliseconds."""
|
||||
|
||||
strategy_selection_ms: float = Field(default=0.0, ge=0.0)
|
||||
budget_allocation_ms: float = Field(default=0.0, ge=0.0)
|
||||
strategy_execution_ms: float = Field(default=0.0, ge=0.0)
|
||||
deduplication_ms: float = Field(default=0.0, ge=0.0)
|
||||
depth_resolution_ms: float = Field(default=0.0, ge=0.0)
|
||||
scoring_ms: float = Field(default=0.0, ge=0.0)
|
||||
packing_ms: float = Field(default=0.0, ge=0.0)
|
||||
ordering_ms: float = Field(default=0.0, ge=0.0)
|
||||
preamble_generation_ms: float = Field(default=0.0, ge=0.0)
|
||||
total_ms: float = Field(default=0.0, ge=0.0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ContextAssemblyPipeline — Full 10-stage mediator
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ContextAssemblyPipeline(ACMSPipeline):
|
||||
"""Full 10-stage ACMS context assembly pipeline mediator.
|
||||
|
||||
Extends ``ACMSPipeline`` with production Phase 1 components:
|
||||
|
||||
- ``ConfidenceWeightedSelector`` (replaces ``DefaultStrategySelector``)
|
||||
- ``ProportionalBudgetAllocator`` (replaces ``DefaultBudgetAllocator``)
|
||||
- ``ParallelStrategyExecutor`` (replaces ``DefaultStrategyExecutor``)
|
||||
|
||||
Adds structured logging with per-stage timing metrics and exposes
|
||||
``StageTimings`` on the most recent assembly result.
|
||||
|
||||
Example::
|
||||
|
||||
pipeline = ContextAssemblyPipeline()
|
||||
payload = pipeline.assemble(
|
||||
plan_id="01ARZ3NDEKTSV4RRFFQ69G5FAV",
|
||||
fragments=[frag1, frag2],
|
||||
budget=ContextBudget(max_tokens=2048),
|
||||
)
|
||||
timings = pipeline.last_timings
|
||||
|
||||
Based on ``docs/specification.md`` ~line 42615.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
default_strategy: str = "relevance",
|
||||
*,
|
||||
settings: Settings | None = None,
|
||||
unit_of_work: UnitOfWork | None = None,
|
||||
# Phase 1 — Strategy Orchestration
|
||||
strategy_selector: StrategySelector | None = None,
|
||||
budget_allocator: BudgetAllocator | None = None,
|
||||
strategy_executor: StrategyExecutor | None = None,
|
||||
# Phase 2 — Fragment Fusion
|
||||
deduplicator: FragmentDeduplicator | None = None,
|
||||
depth_resolver: DetailDepthResolver | None = None,
|
||||
scorer: FragmentScorer | None = None,
|
||||
packer: BudgetPacker | None = None,
|
||||
orderer: FragmentOrderer | None = None,
|
||||
# Phase 3 — Context Finalization
|
||||
preamble_generator: PreambleGenerator | None = None,
|
||||
skeleton_compressor: SkeletonCompressor | None = None,
|
||||
# Phase 1 configuration
|
||||
min_useful_budget: int = 64,
|
||||
executor_timeout: float = 30.0,
|
||||
executor_max_workers: int = 4,
|
||||
) -> None:
|
||||
# Build Phase 1 defaults if not provided
|
||||
effective_selector = strategy_selector or ConfidenceWeightedSelector()
|
||||
effective_allocator = budget_allocator or ProportionalBudgetAllocator(
|
||||
min_useful_budget=min_useful_budget,
|
||||
)
|
||||
effective_executor = strategy_executor or ParallelStrategyExecutor(
|
||||
timeout_seconds=executor_timeout,
|
||||
max_workers=executor_max_workers,
|
||||
)
|
||||
|
||||
super().__init__(
|
||||
default_strategy,
|
||||
settings=settings,
|
||||
unit_of_work=unit_of_work,
|
||||
strategy_selector=effective_selector,
|
||||
budget_allocator=effective_allocator,
|
||||
strategy_executor=effective_executor,
|
||||
deduplicator=deduplicator or DefaultDeduplicator(),
|
||||
depth_resolver=depth_resolver or DefaultDepthResolver(),
|
||||
scorer=scorer or DefaultScorer(),
|
||||
packer=packer or DefaultBudgetPacker(),
|
||||
orderer=orderer or DefaultOrderer(),
|
||||
preamble_generator=preamble_generator or DefaultPreambleGenerator(),
|
||||
skeleton_compressor=skeleton_compressor or DefaultSkeletonCompressor(),
|
||||
)
|
||||
self._last_timings: StageTimings | None = None
|
||||
self._pipeline_logger = logger.bind(service="context_assembly_pipeline")
|
||||
|
||||
@property
|
||||
def last_timings(self) -> StageTimings | None:
|
||||
"""Return timing data from the most recent ``assemble`` call."""
|
||||
return self._last_timings
|
||||
|
||||
def assemble(
|
||||
self,
|
||||
plan_id: str,
|
||||
fragments: Sequence[ContextFragment],
|
||||
budget: ContextBudget,
|
||||
strategy: str | None = None,
|
||||
) -> ContextPayload:
|
||||
"""Assemble context with per-stage timing instrumentation.
|
||||
|
||||
Overrides :meth:`ACMSPipeline.assemble` to add structured logging
|
||||
with per-stage millisecond timings.
|
||||
"""
|
||||
import re
|
||||
|
||||
from cleveragents.domain.models.core.context_fragment import ULID_PATTERN
|
||||
|
||||
if not re.match(ULID_PATTERN, plan_id):
|
||||
msg = f"plan_id must be a valid ULID, got {plan_id!r}"
|
||||
raise ValueError(msg)
|
||||
|
||||
strategy_name = strategy or self._default_strategy
|
||||
if strategy_name not in self._strategies:
|
||||
msg = (
|
||||
f"Unknown strategy {strategy_name!r}. "
|
||||
f"Available: {', '.join(sorted(self._strategies))}"
|
||||
)
|
||||
raise ValueError(msg)
|
||||
|
||||
pipeline_start = time.monotonic()
|
||||
self._pipeline_logger.info(
|
||||
"Pipeline started",
|
||||
plan_id=plan_id,
|
||||
strategy=strategy_name,
|
||||
fragment_count=len(fragments),
|
||||
budget_tokens=budget.available_tokens,
|
||||
)
|
||||
|
||||
# --- Phase 1: Strategy Orchestration ---
|
||||
all_strategies = list(self._strategies.values())
|
||||
resolved = self._strategies[strategy_name]
|
||||
|
||||
t0 = time.monotonic()
|
||||
candidates = self._strategy_selector.select(
|
||||
all_strategies,
|
||||
{"strategy": strategy_name},
|
||||
)
|
||||
candidates = [(s, c) for s, c in candidates if s.name == strategy_name] or [
|
||||
(resolved, 1.0)
|
||||
]
|
||||
selection_ms = (time.monotonic() - t0) * 1000
|
||||
|
||||
t0 = time.monotonic()
|
||||
allocations = self._budget_allocator.allocate(
|
||||
candidates,
|
||||
budget.available_tokens,
|
||||
)
|
||||
allocation_ms = (time.monotonic() - t0) * 1000
|
||||
|
||||
t0 = time.monotonic()
|
||||
ranked = self._strategy_executor.execute(allocations, fragments, budget)
|
||||
execution_ms = (time.monotonic() - t0) * 1000
|
||||
|
||||
# --- Phase 2: Fragment Fusion ---
|
||||
t0 = time.monotonic()
|
||||
fused: Sequence[ContextFragment] = self._deduplicator.deduplicate(ranked)
|
||||
dedup_ms = (time.monotonic() - t0) * 1000
|
||||
|
||||
t0 = time.monotonic()
|
||||
fused = self._depth_resolver.resolve(fused)
|
||||
depth_ms = (time.monotonic() - t0) * 1000
|
||||
|
||||
t0 = time.monotonic()
|
||||
fused = self._scorer.score(fused)
|
||||
score_ms = (time.monotonic() - t0) * 1000
|
||||
|
||||
t0 = time.monotonic()
|
||||
fused = self._packer.pack(fused, budget)
|
||||
pack_ms = (time.monotonic() - t0) * 1000
|
||||
|
||||
t0 = time.monotonic()
|
||||
fused = self._orderer.order(fused)
|
||||
order_ms = (time.monotonic() - t0) * 1000
|
||||
|
||||
# --- Phase 3: Context Finalization ---
|
||||
t0 = time.monotonic()
|
||||
preamble = self._preamble_generator.generate(fused)
|
||||
preamble_ms = (time.monotonic() - t0) * 1000
|
||||
|
||||
total_ms = (time.monotonic() - pipeline_start) * 1000
|
||||
|
||||
self._last_timings = StageTimings(
|
||||
strategy_selection_ms=round(selection_ms, 3),
|
||||
budget_allocation_ms=round(allocation_ms, 3),
|
||||
strategy_execution_ms=round(execution_ms, 3),
|
||||
deduplication_ms=round(dedup_ms, 3),
|
||||
depth_resolution_ms=round(depth_ms, 3),
|
||||
scoring_ms=round(score_ms, 3),
|
||||
packing_ms=round(pack_ms, 3),
|
||||
ordering_ms=round(order_ms, 3),
|
||||
preamble_generation_ms=round(preamble_ms, 3),
|
||||
total_ms=round(total_ms, 3),
|
||||
)
|
||||
|
||||
final_fragments = tuple(fused)
|
||||
total_tokens = sum(f.token_count for f in final_fragments)
|
||||
available = budget.available_tokens
|
||||
budget_used = total_tokens / available if available > 0 else 0.0
|
||||
context_hash = compute_context_hash(final_fragments)
|
||||
provenance_map = build_provenance_map(final_fragments)
|
||||
|
||||
self._pipeline_logger.info(
|
||||
"Pipeline complete",
|
||||
plan_id=plan_id,
|
||||
strategy=strategy_name,
|
||||
fragments_selected=len(final_fragments),
|
||||
total_tokens=total_tokens,
|
||||
budget_used=round(budget_used, 4),
|
||||
total_ms=round(total_ms, 3),
|
||||
)
|
||||
|
||||
return ContextPayload(
|
||||
plan_id=plan_id,
|
||||
fragments=final_fragments,
|
||||
total_tokens=total_tokens,
|
||||
budget=budget,
|
||||
budget_used=round(min(budget_used, 1.0), 4),
|
||||
strategies_used=(strategy_name,),
|
||||
context_hash=context_hash,
|
||||
preamble=preamble,
|
||||
provenance_map=provenance_map,
|
||||
)
|
||||
Reference in New Issue
Block a user