diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/budget.py b/packages/shared-python/shared/services/retrieval/agentic/core/budget.py index ff7c7f2f7..05f99d700 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/core/budget.py +++ b/packages/shared-python/shared/services/retrieval/agentic/core/budget.py @@ -102,7 +102,7 @@ def __init__( total: int, planning_ratio: float, bootstrap: int = 2000, - per_doc_min_share: int = 1500, + per_doc_cap: int = 20000, ) -> None: total = max(int(total), 1) bootstrap = max(0, min(int(bootstrap), total)) @@ -117,10 +117,16 @@ def __init__( "planning": BudgetPool("planning", planning_capacity), "context": BudgetPool("context", context_capacity), } + self._total = ( + self._pools["bootstrap"].capacity + + self._pools["planning"].capacity + + self._pools["context"].capacity + ) self._doc_caps: dict[str, int] = {} self._doc_used: dict[str, int] = {} self._doc_reserved: dict[str, int] = {} - self._per_doc_min_share = max(int(per_doc_min_share), 0) + self._planning_ratio = planning_ratio + self._per_doc_cap = max(int(per_doc_cap), 0) self.total_chunks = 0 self.total_docs = 0 self.explored_chunks = 0 @@ -128,6 +134,11 @@ def __init__( self.trimmed_paths: list[dict[str, Any]] = [] self._overdraft_events: list[dict[str, Any]] = [] + @property + def total(self) -> int: + """Spendable token total across bootstrap/planning/context pools.""" + return self._total + def remaining(self, pool: BudgetPoolName) -> int: return self._pools[pool].remaining @@ -138,25 +149,56 @@ def status(self, pool: BudgetPoolName) -> BudgetStatus: used_pct=pool_state.used_pct, ) - async def allocate_doc_caps(self, doc_chunks: dict[str, int]) -> None: - """Allocate planning soft caps by document chunk counts.""" + async def allocate_doc_caps(self, doc_ids: list[str]) -> None: + """Grow pools and total budget for selected docs with a flat per-doc cap. + + After document selection, planning capacity becomes + ``len(unique_docs) * per_doc_cap`` (not split by chunk count). Context + grows with the same planning/context ratio so evidence budget scales + with the selected set. The ledger ``total`` is raised to match the new + pool capacities. Each document may spend up to ``per_doc_cap``. + """ async with self._lock: self._doc_caps.clear() self._doc_used.clear() self._doc_reserved.clear() - if not doc_chunks: + unique_ids = list(dict.fromkeys(doc_id for doc_id in doc_ids if doc_id)) + if not unique_ids: return - planning_capacity = self._pools["planning"].capacity - total_weight = sum(max(int(count), 1) for count in doc_chunks.values()) - for doc_id, count in doc_chunks.items(): - weight = max(int(count), 1) - weighted = int(planning_capacity * weight / total_weight) - self._doc_caps[doc_id] = min( - planning_capacity, - max(self._per_doc_min_share, weighted), + n_docs = len(unique_ids) + per_doc = self._per_doc_cap + planning_capacity = n_docs * per_doc + if self._planning_ratio <= 0: + context_capacity = self._pools["context"].capacity + elif self._planning_ratio >= 1: + context_capacity = 0 + else: + context_capacity = int( + planning_capacity + * (1.0 - self._planning_ratio) + / self._planning_ratio ) + planning = self._pools["planning"] + context = self._pools["context"] + planning.capacity = max( + planning_capacity, + planning.used + planning.reserved, + ) + context.capacity = max( + context_capacity, + context.used + context.reserved, + ) + self._total = ( + self._pools["bootstrap"].capacity + + planning.capacity + + context.capacity + ) + + for doc_id in unique_ids: + self._doc_caps[doc_id] = per_doc + async def try_reserve( self, pool: BudgetPoolName, @@ -299,6 +341,7 @@ def snapshot(self) -> dict[str, object]: if self._overdraft_events: snapshot["overdraft_events"] = list(self._overdraft_events) snapshot.update({ + "total": self._total, "total_chunks": self.total_chunks, "total_docs": self.total_docs, "explored_chunks": min(self.explored_chunks, self.total_chunks) diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py b/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py index 91d7a9a80..360c4343f 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py +++ b/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py @@ -23,7 +23,7 @@ def build_config_from_env() -> AgentRunConfig: token_budget_total=int(os.environ.get("RETRIEVAL_AGENTIC_TOKEN_BUDGET_TOTAL", "40000")), planning_ratio=float(os.environ.get("RETRIEVAL_AGENTIC_PLANNING_RATIO", "0.5")), bootstrap_budget=int(os.environ.get("RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET", "2000")), - per_doc_min_share=int(os.environ.get("RETRIEVAL_AGENTIC_PER_DOC_MIN_SHARE", "1500")), + per_doc_cap=int(os.environ.get("RETRIEVAL_AGENTIC_PER_DOC_CAP", "20000")), inventory_aware=os.environ.get("RETRIEVAL_AGENTIC_INVENTORY_AWARE", "true") == "true", ) diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/types.py b/packages/shared-python/shared/services/retrieval/agentic/core/types.py index a3bb1293b..194f864e6 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/core/types.py +++ b/packages/shared-python/shared/services/retrieval/agentic/core/types.py @@ -21,7 +21,7 @@ class AgentRunConfig: token_budget_total: int = 40000 planning_ratio: float = 0.5 bootstrap_budget: int = 2000 - per_doc_min_share: int = 1500 + per_doc_cap: int = 20000 inventory_aware: bool = True diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py index 32a604966..ced49e0cd 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py @@ -99,9 +99,9 @@ async def run( total=config.token_budget_total, planning_ratio=config.planning_ratio, bootstrap=config.bootstrap_budget, - per_doc_min_share=config.per_doc_min_share, + per_doc_cap=config.per_doc_cap, ) - total_chunks, total_docs, chunks_count_by_doc = await _load_budget_inventory( + total_chunks, total_docs, _chunks_count_by_doc = await _load_budget_inventory( db, user_id=user_id, namespace=namespace, @@ -267,10 +267,13 @@ async def run( if state.ledger is not None: - await state.ledger.allocate_doc_caps({ - doc.document_id: chunks_count_by_doc.get(doc.document_id, 1) - for doc in state.selected_docs - }) + await state.ledger.allocate_doc_caps( + [doc.document_id for doc in state.selected_docs] + ) + logger.info( + f'agentic: doc caps allocated for {len(state.selected_docs)} docs; ' + f'token_budget_total={state.ledger.total}' + ) # Phase 2 + 3: navigate once, then render evidence for downstream agents. evidence_text = '' diff --git a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py index aa4db437a..ddc268c5e 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py @@ -117,7 +117,7 @@ async def run_request( total=config.planner_budget, planning_ratio=0.0, bootstrap=config.planner_budget, - per_doc_min_share=0, + per_doc_cap=0, ) db_factory = self._get_db_factory() async with db_factory() as inventory_db: diff --git a/packages/shared-python/shared/services/retrieval/workflow/wallet.py b/packages/shared-python/shared/services/retrieval/workflow/wallet.py index da9d67eac..19216bb7e 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/wallet.py +++ b/packages/shared-python/shared/services/retrieval/workflow/wallet.py @@ -39,8 +39,8 @@ class BudgetWallet: bootstrap_budget: int = field( default_factory=lambda: _env_int('RETRIEVAL_AGENTIC_BOOTSTRAP_BUDGET', 2000) ) - per_doc_min_share: int = field( - default_factory=lambda: _env_int('RETRIEVAL_AGENTIC_PER_DOC_MIN_SHARE', 1500) + per_doc_cap: int = field( + default_factory=lambda: _env_int('RETRIEVAL_AGENTIC_PER_DOC_CAP', 20000) ) _allocations: dict[str, int] = field(default_factory=dict, init=False) _reclaimed: dict[str, int] = field(default_factory=dict, init=False) @@ -89,8 +89,14 @@ async def allocate(self, plan: QueryPlan) -> dict[str, BudgetLedger]: return dict(self._ledgers) async def reclaim(self, step_id: str, ledger: BudgetLedger) -> None: - """Record unused capacity after a step completes.""" + """Record unused capacity after a step completes. + + If the step ledger grew its total after document selection + (``allocate_doc_caps``), raise this wallet's hard total and the step + allocation to match so accounting stays consistent. + """ async with self._lock: + self._sync_capacity_growth(step_id, ledger) allocated = self._allocations.get(step_id, 0) used = self._ledger_used(ledger) self._reclaimed[step_id] = max(allocated - used, 0) @@ -99,12 +105,14 @@ def total_used(self) -> int: return sum(self._ledger_used(ledger) for ledger in self._ledgers.values()) def snapshot(self) -> dict[str, object]: + allocations, total = self._live_allocations() + used = self.total_used() return { - "total": self.total, - "allocated": sum(self._allocations.values()), - "used": self.total_used(), - "remaining": max(self.total - self.total_used(), 0), - "allocations": dict(self._allocations), + "total": total, + "allocated": sum(allocations.values()), + "used": used, + "remaining": max(total - used, 0), + "allocations": allocations, "reclaimed": dict(self._reclaimed), "steps": { step_id: ledger.snapshot() @@ -112,6 +120,25 @@ def snapshot(self) -> dict[str, object]: }, } + def _sync_capacity_growth(self, step_id: str, ledger: BudgetLedger) -> None: + capacity = max(int(ledger.total), 0) + previous = self._allocations.get(step_id, 0) + if capacity > previous: + self.total += capacity - previous + self._allocations[step_id] = capacity + + def _live_allocations(self) -> tuple[dict[str, int], int]: + """Return allocations/total including any post-selection ledger growth.""" + allocations = dict(self._allocations) + total = self.total + for step_id, ledger in self._ledgers.items(): + capacity = max(int(ledger.total), 0) + previous = allocations.get(step_id, 0) + if capacity > previous: + total += capacity - previous + allocations[step_id] = capacity + return allocations, total + def _requested_for_step(self, step: PlannedStep) -> int: del step return max(self.per_retrieve_step_default, _RETRIEVE_FLOOR) @@ -122,7 +149,7 @@ def _new_ledger(self, step: PlannedStep, total: int) -> BudgetLedger: total=max(total, 1), planning_ratio=self.planning_ratio, bootstrap=min(self.bootstrap_budget, max(total, 1)), - per_doc_min_share=self.per_doc_min_share, + per_doc_cap=self.per_doc_cap, ) @staticmethod