feat(api): opt-in per-bank attribution for remote reranker requests - #2
Open
chethanuk wants to merge 8 commits into
Open
feat(api): opt-in per-bank attribution for remote reranker requests#2chethanuk wants to merge 8 commits into
chethanuk wants to merge 8 commits into
Conversation
|
CodeAnt AI is reviewing your PR. |
Thanks for using CodeAnt! 🎉We're free for open-source projects. if you're enjoying it, help us grow by sharing. Share on X · |
📝 WalkthroughWalkthroughThe changes add preflight validation for bank operations and configuration updates, separate configuration validation from persistence, and introduce optional per-request bank-id headers for Cohere-compatible remote reranking with accompanying configuration, documentation, and tests. ChangesBank configuration validation
Reranker bank attribution
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant RerankCaller
participant BankContext
participant CohereCompatibleRerankClient
participant RemoteRerankEndpoint
RerankCaller->>BankContext: bind current bank id
CohereCompatibleRerankClient->>BankContext: read attribution context
BankContext-->>CohereCompatibleRerankClient: X-Hindsight-Bank-Id header
CohereCompatibleRerankClient->>RemoteRerankEndpoint: POST /rerank with headers
sequenceDiagram
participant BankUpdateAPI
participant ConfigResolver
participant BankStore
BankUpdateAPI->>ConfigResolver: validate bank operations and updates
ConfigResolver-->>BankUpdateAPI: normalized updates
BankUpdateAPI->>BankStore: persist validated updates
BankStore-->>BankUpdateAPI: return latest stored configuration
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
codeant-ai
Bot
added
the
size:L
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an opt-in configuration (HINDSIGHT_API_RERANKER_SEND_BANK_AS_HEADER) to tag outbound Cohere-compatible remote rerank requests with X-Hindsight-Bank-Id for per-bank cost attribution. It also refactors bank configuration updates to pre-validate payloads and permissions before executing any side effects across several API endpoints and MCP tools. The review feedback highlights a redundancy where validate_bank_config_updates is called first, but the subsequent call to update_bank_config runs the exact same validation again, and suggests refactoring to avoid this duplicate processing.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Comment on lines
+438
to
+455
| async def update_bank_config( | ||
| self, bank_id: str, updates: dict[str, Any], context: RequestContext | None = None | ||
| ) -> None: | ||
| """ | ||
| Update bank configuration overrides (with permission checking). | ||
|
|
||
| Args: | ||
| bank_id: Bank identifier | ||
| updates: Dict of config field names to new values. | ||
| Keys can be in env var format (HINDSIGHT_API_LLM_PROVIDER) | ||
| or Python field format (llm_provider). | ||
| Only configurable fields are allowed. | ||
| context: Request context for permission checking | ||
|
|
||
| Raises: | ||
| ValueError: If attempting to override invalid/disallowed fields | ||
| """ | ||
| normalized_updates = await self.validate_bank_config_updates(bank_id, updates, context) |
There was a problem hiding this comment.
With the introduction of validate_bank_config_updates, several endpoints (e.g., api_create_or_update_bank, api_update_bank, api_import_bank_template, api_update_bank_config, and _do_update_bank in mcp_tools.py) now call validate_bank_config_updates first to validate the payload before any side effects, and then call update_bank_config which runs the exact same validation again.
To avoid this redundant validation and potential extra database/cache lookups, consider splitting the persistence logic out of update_bank_config into a separate method (e.g., _persist_bank_config_updates), or allowing update_bank_config to accept pre-validated/normalized updates.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@hindsight-api-slim/hindsight_api/api/http.py`:
- Around line 5661-5663: Wrap the validate_bank_config_updates call in both bank
update handlers with ValueError handling and map validation or permission
failures to HTTP 400, matching the existing config and import endpoint behavior.
Apply this at hindsight-api-slim/hindsight_api/api/http.py lines 5661-5663 and
5730-5732, preserving validation before any side effects.
- Around line 5861-5869: Extend the pre-validation block around
_verify_operation to also validate the update operations used by the import flow
for existing mental models and directives. When body.mental_models or
body.directives is present, check the corresponding update permission in
addition to the existing create permission checks, so all possible mutations are
authorized before changes begin.
In `@hindsight-api-slim/hindsight_api/config_resolver.py`:
- Line 455: Update the compound bank-update flow around
validate_bank_config_updates to persist its returned normalized_updates directly
instead of invoking asynchronous validation again. Add or reuse a
persistence-only path that accepts the prevalidated result, and ensure the
bank-name mutation and normalized configuration update commit atomically so a
persistence failure cannot leave a partial name change.
- Around line 466-468: Wrap the bank initialization, row fetch, overlay
computation, and final update block in an explicit transaction using
self._backend.transaction() or the acquired connection’s transaction context.
Ensure ensure_bank_exists() and the UPDATE execute within that same transaction,
and revise the nearby comment to accurately describe the consolidated steps.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: d096bd3f-4f91-43e3-ba42-efe31e458188
📒 Files selected for processing (11)
.env.examplehindsight-api-slim/hindsight_api/api/http.pyhindsight-api-slim/hindsight_api/config.pyhindsight-api-slim/hindsight_api/config_resolver.pyhindsight-api-slim/hindsight_api/engine/bank_attribution.pyhindsight-api-slim/hindsight_api/engine/cross_encoder.pyhindsight-api-slim/hindsight_api/mcp_tools.pyhindsight-api-slim/tests/test_bank_attribution_config.pyhindsight-api-slim/tests/test_reranker_bank_attribution.pyhindsight-docs/docs/developer/configuration.mdhindsight-embed/hindsight_embed/env.example
Comment on lines
+5661
to
+5663
| if config_updates: | ||
| # Validate config payload and permissions BEFORE any side effects | ||
| await app.state.memory._config_resolver.validate_bank_config_updates(bank_id, config_updates, request_context) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Map preflight configuration failures to a client error.
Both handlers now invoke validate_bank_config_updates, but neither catches its ValueError, so invalid or permission-denied config requests return HTTP 500.
hindsight-api-slim/hindsight_api/api/http.py#L5661-L5663: catchValueErrorand return HTTP 400, consistent with the config and import endpoints.hindsight-api-slim/hindsight_api/api/http.py#L5730-L5732: apply the same mapping to PATCH bank updates.
📍 Affects 1 file
hindsight-api-slim/hindsight_api/api/http.py#L5661-L5663(this comment)hindsight-api-slim/hindsight_api/api/http.py#L5730-L5732
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@hindsight-api-slim/hindsight_api/api/http.py` around lines 5661 - 5663, Wrap
the validate_bank_config_updates call in both bank update handlers with
ValueError handling and map validation or permission failures to HTTP 400,
matching the existing config and import endpoint behavior. Apply this at
hindsight-api-slim/hindsight_api/api/http.py lines 5661-5663 and 5730-5732,
preserving validation before any side effects.
Comment on lines
+5861
to
+5869
| if hasattr(app.state.memory, "_verify_operation"): | ||
| # Pre-validate all requested operations before making any changes | ||
| await app.state.memory._verify_operation(bank_id, "get_bank_profile", request_context) | ||
| if body.bank: | ||
| await app.state.memory._verify_operation(bank_id, "update_bank_config", request_context) | ||
| if body.mental_models: | ||
| await app.state.memory._verify_operation(bank_id, "create_mental_model", request_context) | ||
| if body.directives: | ||
| await app.state.memory._verify_operation(bank_id, "create_directive", request_context) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n -C5 'def apply_bank_template_manifest|async def apply_bank_template_manifest' .
rg -n -C4 '\b(create|update)_(mental_model|directive)\s*\(' hindsight-api-slim/hindsight_apiRepository: chethanuk/hindsight
Length of output: 23787
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate permission/operation validation logic and the relevant operation names.
rg -n -C3 '_verify_operation|create_mental_model|update_mental_model|create_directive|update_directive|OperationValidationError|permission' hindsight-api-slim/hindsight_api | sed -n '1,260p'Repository: chethanuk/hindsight
Length of output: 21951
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the route area around the preflight checks and the import helper.
sed -n '2440,2565p' hindsight-api-slim/hindsight_api/api/http.py
printf '\n----\n'
sed -n '5840,5875p' hindsight-api-slim/hindsight_api/api/http.pyRepository: chethanuk/hindsight
Length of output: 7376
Preflight the update paths too. The import flow updates existing mental models and directives, but this guard only checks create_mental_model and create_directive, so a request can partially mutate the bank before an update permission failure aborts the rest.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@hindsight-api-slim/hindsight_api/api/http.py` around lines 5861 - 5869,
Extend the pre-validation block around _verify_operation to also validate the
update operations used by the import flow for existing mental models and
directives. When body.mental_models or body.directives is present, check the
corresponding update permission in addition to the existing create permission
checks, so all possible mutations are authorized before changes begin.
| Raises: | ||
| ValueError: If attempting to override invalid/disallowed fields | ||
| """ | ||
| normalized_updates = await self.validate_bank_config_updates(bank_id, updates, context) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Persist the prevalidated result without validating again.
HTTP and MCP callers validate first, mutate the bank name, then this method repeats asynchronous permission/config validation. If the second pass rejects, the name remains changed and the compound update is only partially applied. Add a persistence path accepting the returned normalized_updates, or make the complete update atomic.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@hindsight-api-slim/hindsight_api/config_resolver.py` at line 455, Update the
compound bank-update flow around validate_bank_config_updates to persist its
returned normalized_updates directly instead of invoking asynchronous validation
again. Add or reuse a persistence-only path that accepts the prevalidated
result, and ensure the bank-name mutation and normalized configuration update
commit atomically so a persistence failure cannot leave a partial name change.
Comment on lines
+466
to
+468
| # We explicitly execute the initial ensure_bank_exists call in a transaction to safely handle the | ||
| # ON CONFLICT DO NOTHING insertion, then fetch the updated row and compute the overlay, | ||
| # then finally update in the same transaction to prevent race conditions. |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
fd -t f 'database|backend|fact_storage' hindsight-api-slim/hindsight_api |
xargs -r rg -n -C4 'def acquire|async def acquire|ensure_bank_exists|conn\.transaction'Repository: chethanuk/hindsight
Length of output: 1003
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the relevant files and inspect the surrounding implementations.
printf '\n## config_resolver.py around referenced lines\n'
sed -n '430,500p' hindsight-api-slim/hindsight_api/config_resolver.py
printf '\n## fact_storage.py around ensure_bank_exists\n'
sed -n '135,240p' hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py
printf '\n## Search for DatabaseBackend acquire implementation\n'
rg -n -C4 'class DatabaseBackend|def acquire|async def acquire|transaction\(' hindsight-api-slim/hindsight_apiRepository: chethanuk/hindsight
Length of output: 50375
Wrap this block in a transaction
self._backend.acquire() only opens a connection; it does not start a transaction. ensure_bank_exists() and the UPDATE can commit separately, so use async with self._backend.transaction() as conn: (or async with conn.transaction():) here and trim the comment to match the actual steps.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@hindsight-api-slim/hindsight_api/config_resolver.py` around lines 466 - 468,
Wrap the bank initialization, row fetch, overlay computation, and final update
block in an explicit transaction using self._backend.transaction() or the
acquired connection’s transaction context. Ensure ensure_bank_exists() and the
UPDATE execute within that same transaction, and revise the nearby comment to
accurately describe the consolidated steps.
![codeant-ai[bot]](https://avatars.githubusercontent.com/in/646884?s=60&v=4){kind=small}
Comment on lines
+466
to
+468
| # We explicitly execute the initial ensure_bank_exists call in a transaction to safely handle the | ||
| # ON CONFLICT DO NOTHING insertion, then fetch the updated row and compute the overlay, | ||
| # then finally update in the same transaction to prevent race conditions. |
There was a problem hiding this comment.
Suggestion: The new comment states this block runs in a single transaction to prevent race conditions, but no transaction is actually opened here (only a pooled connection is acquired). That mismatch is dangerous because readers will assume atomicity that is not guaranteed; either wrap the sequence in an explicit transaction (conn.transaction()) or correct the comment to avoid documenting behavior that the code does not provide. [comment mismatch]
Severity Level: Major ⚠️
- ⚠️ Bank config PATCH handler not truly transactional as documented.
- ⚠️ Concurrent config updates may interleave despite transaction comment.
- ⚠️ Future maintainers may rely on nonexistent race protection.Steps of Reproduction ✅
1. Open `hindsight-api-slim/hindsight_api/config_resolver.py` and locate
`update_bank_config()`; see the comment at lines 466-468 stating the sequence runs "in a
transaction" to prevent race conditions.
2. Inspect the context in `update_bank_config()`: the comment sits above `async with
self._backend.acquire() as conn:` and the calls to `ensure_bank_exists(conn, bank_id,
ops=self._backend.ops)` and `conn.execute(...)` (lines 465-470).
3. Check the backend implementation in
`hindsight-api-slim/hindsight_api/engine/db/base.py:11-33`, where `acquire()` is defined
as an async context manager that "Acquire[s] a connection from the pool" and yields a
`DatabaseConnection` wrapper without starting a transaction.
4. Verify the concrete PostgreSQL backend in
`hindsight-api-slim/hindsight_api/engine/db/postgresql.py:4-15`: `acquire()` (lines 4-8)
wraps `pool.acquire()` only, while `transaction()` (lines 10-15) explicitly uses `async
with conn.transaction()` to start a database transaction, demonstrating that
`update_bank_config()` currently uses a non-transactional acquire despite the
transactional comment.
5. Trace a real caller: the bank config update HTTP handler in
`hindsight-api-slim/hindsight_api/api/http.py:6328-21` calls `await
app.state.memory._config_resolver.update_bank_config(bank_id, request.updates,
request_context)` whenever the bank config PATCH endpoint is hit, meaning each live config
update relies on this non-transactional sequence even though the comment claims
transaction-based race protection.
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** hindsight-api-slim/hindsight_api/config_resolver.py
**Line:** 466:468
**Comment:**
*Comment Mismatch: The new comment states this block runs in a single transaction to prevent race conditions, but no transaction is actually opened here (only a pooled connection is acquired). That mismatch is dangerous because readers will assume atomicity that is not guaranteed; either wrap the sequence in an explicit transaction (`conn.transaction()`) or correct the comment to avoid documenting behavior that the code does not provide.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
Comment on lines
+3240
to
+3241
| if mission is not None and "reflect_mission" not in effective_config: | ||
| effective_config["reflect_mission"] = mission |
There was a problem hiding this comment.
Suggestion: The precedence check only looks for the Python-style key reflect_mission, so if callers provide the equivalent env-style key (for example HINDSIGHT_API_REFLECT_MISSION) and also pass mission, the alias value is appended and then wins during normalization. This silently overrides the explicit config_updates value with the deprecated alias; normalize keys first (or check both key forms) before applying the mission fallback. [incorrect condition logic]
Severity Level: Major ⚠️
- ⚠️ MCP update_bank tool may ignore env-style reflect_mission.
- ⚠️ Mixed env/Python reflect_mission overrides resolve to deprecated alias.
- ⚠️ Bank reflect behavior may not match explicit config_updates.Steps of Reproduction ✅
1. Create an MCP server using the helper `_make_mcp_server` defined in
`hindsight-api-slim/tests/test_mcp_tools.py:979-989`, which calls `register_mcp_tools(mcp,
mock_memory, config)` and registers the `update_bank` tool implemented via
`_do_update_bank` in `hindsight-api-slim/hindsight_api/mcp_tools.py:3221-3320`.
2. Invoke the registered MCP tool function for `update_bank` (obtained via
`_tools(mcp)["update_bank"].fn` from `tests/test_mcp_tools.py:500-505`) with both an
env-style key and the deprecated alias, for example:
`await _tools(mcp)["update_bank"].fn(mission="alias-mission",
config_updates={"HINDSIGHT_API_REFLECT_MISSION": "explicit-mission"})`.
3. Inside `_do_update_bank` (`hindsight_api/mcp_tools.py:3221-3242`), `effective_config`
is initialized from `config_updates`, so it contains the key
`"HINDSIGHT_API_REFLECT_MISSION"` but not `"reflect_mission"`. Because the condition at
line 3240 only checks for `"reflect_mission"` in `effective_config`, it evaluates true and
line 3241 writes `effective_config["reflect_mission"] = mission`, adding a second key for
the same conceptual setting.
4. `_do_update_bank` then calls
`memory._config_resolver.validate_bank_config_updates(target_bank, effective_config,
request_context)` at lines 3253-3255, which in turn invokes `normalize_config_dict` in
`hindsight-api-slim/hindsight_api/config_resolver.py:15-24` and
`hindsight-api-slim/hindsight_api/config.py:86-107`. Both
`"HINDSIGHT_API_REFLECT_MISSION"` and `"reflect_mission"` normalize to
`"reflect_mission"`, and due to dict overwrite semantics, the later alias-derived value
from `mission` overwrites the earlier env-style `config_updates` value, so the final
update uses `"alias-mission"` instead of the explicitly provided `"explicit-mission"`.
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** hindsight-api-slim/hindsight_api/mcp_tools.py
**Line:** 3240:3241
**Comment:**
*Incorrect Condition Logic: The precedence check only looks for the Python-style key `reflect_mission`, so if callers provide the equivalent env-style key (for example `HINDSIGHT_API_REFLECT_MISSION`) and also pass `mission`, the alias value is appended and then wins during normalization. This silently overrides the explicit `config_updates` value with the deprecated alias; normalize keys first (or check both key forms) before applying the `mission` fallback.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
Comment on lines
+3253
to
+3254
| if effective_config: | ||
| await memory._config_resolver.validate_bank_config_updates(target_bank, effective_config, request_context) |
There was a problem hiding this comment.
Suggestion: This introduces a hard dependency on a private async method on memory._config_resolver; implementations/mocks that only provide update_bank_config will now fail with AttributeError or TypeError when this await runs. Use a capability check/fallback (or rely on update_bank_config’s built-in validation) so _do_update_bank does not require an extra private API contract. [api mismatch]
Severity Level: Major ⚠️
- ❌ MCP update_bank tests crash with AttributeError on validation call.
- ⚠️ Custom MemoryEngine mocks without validation method break MCP tools.
- ⚠️ MCP integration tightly couples to private ConfigResolver API.Steps of Reproduction ✅
1. The `mock_memory` fixture in `hindsight-api-slim/tests/test_mcp_tools.py:124-203`
constructs a `MagicMock` to stand in for the memory engine; it assigns
`memory._config_resolver = MagicMock()` and only defines
`memory._config_resolver.update_bank_config = AsyncMock()`, without providing any
`validate_bank_config_updates` attribute.
2. `_make_mcp_server` in `tests/test_mcp_tools.py:979-989` calls `register_mcp_tools(mcp,
mock_memory, config)`, which wires the `update_bank` MCP tool to the shared implementation
`_do_update_bank` defined in `hindsight-api-slim/hindsight_api/mcp_tools.py:3221-3269`.
3. Tests such as `test_update_bank_with_bank_id_override` in
`tests/test_mcp_tools.py:1720-1729` invoke the tool via
`_tools(mcp)["update_bank"].fn(config_updates={"reflect_mission": "Test"},
bank_id="other-bank")`, causing `_do_update_bank` to execute with a non-empty
`effective_config`.
4. When execution reaches the new validation block at
`hindsight-api-slim/hindsight_api/mcp_tools.py:3252-3255`, the condition `if
effective_config:` is true and Python evaluates `await
memory._config_resolver.validate_bank_config_updates(...)`. Because the mock
`_config_resolver` from step 1 has no `validate_bank_config_updates` attribute, this
raises an `AttributeError` at runtime, breaking all such test paths (and any real caller
whose engine mock or alternative implementation only provides
`_config_resolver.update_bank_config` but not `validate_bank_config_updates`).
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** hindsight-api-slim/hindsight_api/mcp_tools.py
**Line:** 3253:3254
**Comment:**
*Api Mismatch: This introduces a hard dependency on a private async method on `memory._config_resolver`; implementations/mocks that only provide `update_bank_config` will now fail with `AttributeError` or `TypeError` when this `await` runs. Use a capability check/fallback (or rely on `update_bank_config`’s built-in validation) so `_do_update_bank` does not require an extra private API contract.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
|
CodeAnt AI finished reviewing your PR. |
…rize-io#2759) tool_expand zipped memory_ids against valid_uuids, which only collects the ids that parsed as UUIDs. One invalid id shifts every later pair by one, so a memory comes back stamped with a different memory's id, and zip truncates the tail so the last requested id gets no entry at all. Key each id to its own UUID and iterate memory_ids directly, so an invalid id can only affect its own entry.
…unters (vectorize-io#2758) The OpenAI-compatible provider extracts cached_tokens and thoughts_tokens on both call paths and hands them to TokenUsage, but never passes them to metrics.record_llm_call, which accepts and buckets both. Two separate effects: - Reasoning tokens reach no counter at all. vectorize-io#2378 made output_tokens visible-only by subtracting thoughts_tokens directly above the record_llm_call, so the reasoning half of the billed output was removed from the metrics path rather than moved onto llm_tokens_thoughts. Before vectorize-io#2378 those tokens were still counted inside output_tokens. - cached_input_tokens has read 0 for every OpenAI-compatible provider since the counter was added; only gemini_llm passes it. Pass both kwargs at the two call sites that parse a usage object. The fallback path (no usage) and the Ollama native path (no reasoning or cached fields) are unchanged. Invariant: recorded output_tokens + recorded thoughts_tokens equals the provider's completion_tokens, so every billed token lands on exactly one counter. The new tests assert on the collector itself; the existing ones patch it without asserting, which is why this went unnoticed.
…io#2764) Bind Reflect to the existing per-bank ContextVar so its tool loop and final synthesis preserve provider cost attribution. Replace two direct ContextVar implementation tests with one integrated Reflect binding/reset regression.
…t MagicMock crash (vectorize-io#2776) test_null_content_recovers_on_retry failed intermittently on the test-api shard with: hindsight_api/metrics.py:591: TypeError: '>' not supported between instances of 'MagicMock' and 'int' (if cached_input_tokens > 0) The mock in _make_chat_response set completion_tokens_details but not the cached-token fields, so the cached-token extraction (openai_compatible_llm.py:948 `response_usage.cached_tokens`, and the prompt_tokens_details path) read an auto-MagicMock and passed it to the metrics recorder. It only surfaced when the metrics path actually ran — which depends on telemetry state that leaks across pytest-xdist workers — so it presented as an intermittent, co-scheduling-dependent failure rather than a deterministic one. Set usage.cached_tokens = 0 and usage.prompt_tokens_details = None so both extraction paths yield int 0. Verified: both tests pass and both paths return int 0 (no MagicMock reaches the `> 0` comparison).
…y) (vectorize-io#2727) Clears the pydantic-settings Dependabot alert across all affected manifests plus the three high-severity alerts in the root lock. pydantic-settings 2.12.0/2.14.0/2.14.1 -> 2.14.2 GHSA-4xgf-cpjx-pc3j transformers 5.3.0 -> 5.12.1 GHSA-fgcw-684q-jj6r soupsieve 2.8 -> 2.8.4 GHSA-2wc2-fm75-p42x GHSA-836r-79rf-4m37 pydantic-settings is transitive everywhere (no direct declaration), so the locks are the only lever. crewai is deliberately left at 2.10.1: the advisory's range is >=2.12.0,<2.14.2 and NestedSecretsSettingsSource did not exist in 2.10.x, so it is unaffected. transformers is a direct dep, and hindsight-api is published, so the declared floor -- not our lock -- is what protects installers of the local-ml/local-onnx extras. The old >=4.53.0 floor resolved to 4.57.6 (vulnerable) under any downstream cap of transformers<5, so raise it to the advisory's first patched version. Note this now fails resolution for consumers pinned below transformers 5 rather than silently installing a vulnerable build. The >=4.53.0 floor was already unreachable in practice: 4.53.0 requires tokenizers<0.22, which our own cap excludes. The tokenizers<=0.23.0 cap is kept. vectorize-io#2055 was caused by transformers declaring a wider tokenizers range in metadata than its import-time check enforces, and the cap is what blocks that; the comment now records this so it does not read as removable. Root uv.lock is reformatted from lock revision 1 to 3 because uv rewrites in its current format whenever it writes. The other 32 locks in the repo are already revision 3 and CI's setup-uv is unpinned, so this aligns root rather than drifting it. Only 3 versions actually change. Verified: local-ml sync resolves tokenizers 0.22.2 under transformers 5.12.1; LocalSTEmbeddings and LocalSTCrossEncoder both initialize and run (the vectorize-io#2055 import path). Lint passes.
chethanuk
force-pushed
the
fix/issue-2729-featapi-optin-perbank-attribution-for-remote
branch
from
July 17, 2026 15:37
d5d761e to
83e4c46
Compare
|
CodeAnt AI is running Incremental review |
Thanks for using CodeAnt! 🎉We're free for open-source projects. if you're enjoying it, help us grow by sharing. Share on X · |
codeant-ai
Bot
added
size:XL
|
CodeAnt AI Incremental review completed. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
User description
Summary
Opt-in per-bank cost attribution for remote rerank requests. When
HINDSIGHT_API_RERANKER_SEND_BANK_AS_HEADER=true, outbound rerankPOSTs carryX-Hindsight-Bank-Id: <bank_id>. Default isfalse— a byte-identical no-op.Covers: the Cohere-compatible HTTP client (Cohere via base URL, OpenRouter, ZeroEntropy, SiliconFlow, Alibaba) and the LiteLLM proxy reranker (
LiteLLMCrossEncoder, same per-request pattern).Fixes vectorize-io#2729
Root cause
_CohereCompatibleRerankClient.predictposted with no per-request headers. The process-wide client only carriesAuthorization/Content-Typeand is shared across banks, so a client-default header could pin bank A's id onto bank B's requests. The existingapply_bank_attribution()body mutator can't help — the Cohere rerank wire format has no OpenAI-styleuserfield.sequenceDiagram participant R as recall / reflect participant CV as _current_bank_id (contextvar) participant C as shared rerank client<br/>(one per process, all banks) participant P as remote provider R->>CV: bind bank_id R->>C: predict(query, docs) C->>C: bank_attribution_headers()<br/>flag on + bank bound → {X-Hindsight-Bank-Id} C->>P: POST /rerank (per-request headers,<br/>never client defaults) Note over C,P: bank A and bank B requests on the same client<br/>each carry their own id — no bleedDesign notes
X-Bank-Idspelling for MCP surfaces is left to maintainers — trivial rename if preferred.reflectcurrently run with no bank bound; upstream fix(reflect): preserve bank attribution in provider calls vectorize-io/hindsight#2764 fixes that binding orthogonally, and this feature picks it up automatically once merged.How this differs from vectorize-io#2740
Strictly additive and scoped to the issue's stated minimum plus the safe LiteLLM proxy path: no existing attribution/LiteLLM tests are deleted, the LiteLLM legacy response fallback is untouched, and the unrelated reflect-binding change stays in vectorize-io#2764. The LiteLLM SDK path is deliberately left out: passing
headers=intolitellm.arerankunconditionally is unverified against litellm's rerank signature (it may break SDK reranks even with the flag off). Tests here exercise the real env flag + contextvar end-to-end throughhttpx.MockTransport— including absent-by-default, absent-when-no-bank, auth-header survival, and two-bank per-request isolation — rather than patching the header helper.Out of scope (follow-ups)
Native Cohere SDK, LiteLLM SDK (pending kwarg verification), Google, TEI (self-hosted; attribution value low), local/FlashRank/JinaMLX/RRF (no remote spend to attribute).
How to test
Test evidence
Run on the de-stacked, squashed branch (single commit on upstream
main):hindsight-embed/tests/test_env_template.pyruff check/ruff format --checkon touched filesty checkonbank_attribution.py+cross_encoder.pygenerate-docs-skill.shidempotencyCodeAnt-AI Description
Add opt-in bank attribution for remote rerank calls and fix token and memory reporting
What Changed
X-Hindsight-Bank-Idwhen enabled, so shared reranker gateways can track spend per bankImpact
✅ Per-bank rerank billing✅ Fewer missing token usage counts✅ Correct memory lookup results💡 Usage Guide
Checking Your Pull Request
Every time you make a pull request, our system automatically looks through it. We check for security issues, mistakes in how you're setting up your infrastructure, and common code problems. We do this to make sure your changes are solid and won't cause any trouble later.
Talking to CodeAnt AI
Got a question or need a hand with something in your pull request? You can easily get in touch with CodeAnt AI right here. Just type the following in a comment on your pull request, and replace "Your question here" with whatever you want to ask:
This lets you have a chat with CodeAnt AI about your pull request, making it easier to understand and improve your code.
Example
Preserve Org Learnings with CodeAnt
You can record team preferences so CodeAnt AI applies them in future reviews. Reply directly to the specific CodeAnt AI suggestion (in the same thread) and replace "Your feedback here" with your input:
This helps CodeAnt AI learn and adapt to your team's coding style and standards.
Example
Retrigger review
Ask CodeAnt AI to review the PR again, by typing:
Check Your Repository Health
To analyze the health of your code repository, visit our dashboard at https://app.codeant.ai. This tool helps you identify potential issues and areas for improvement in your codebase, ensuring your repository maintains high standards of code health.