diff --git a/.env.example b/.env.example index df447edcb7..c8dc867de4 100644 --- a/.env.example +++ b/.env.example @@ -165,8 +165,11 @@ HINDSIGHT_API_LLM_MODEL=gpt-4o-mini # two Codex members fails over between two ChatGPT accounts instead of retrying one. # HINDSIGHT_API_LLM_1_PROVIDER=openai-codex # HINDSIGHT_API_LLM_1_CODEX_HOME=/var/lib/hindsight/codex-b -# Strategy JSON: {"mode": "failover"} or {"mode": "round-robin"}. +# Strategy JSON: {"mode": "failover"}, {"mode": "round-robin"}, or +# {"mode": "metadata", "routes": [{"key": "tags", "value": "sensitive", "member": 1}]}. # Round-robin accepts optional positive-int "weights" (one per member, primary first). +# Metadata routes are strict: all matches must select one member. An unmatched +# request uses member 0 (the primary); a match never falls across lanes. # HINDSIGHT_API_LLM_STRATEGY={"mode": "failover"} # API Configuration (Optional) diff --git a/hindsight-api-slim/hindsight_api/api/http.py b/hindsight-api-slim/hindsight_api/api/http.py index 8baf6b832c..92b3514a14 100644 --- a/hindsight-api-slim/hindsight_api/api/http.py +++ b/hindsight-api-slim/hindsight_api/api/http.py @@ -8310,14 +8310,14 @@ async def api_retain( content_dict["event_date"] = item.timestamp if item.context: content_dict["context"] = item.context - if item.metadata: + if item.metadata is not None: content_dict["metadata"] = item.metadata if item.document_id: content_dict["document_id"] = item.document_id if item.entities: content_dict["entities"] = [{"text": e.text, "type": e.type or "CONCEPT"} for e in item.entities] content_dict["resolve_entities"] = item.resolve_entities - if item.tags: + if item.tags is not None: content_dict["tags"] = item.tags if item.observation_scopes is not None: content_dict["observation_scopes"] = item.observation_scopes diff --git a/hindsight-api-slim/hindsight_api/config.py b/hindsight-api-slim/hindsight_api/config.py index f1618cfeb6..65b1b84055 100644 --- a/hindsight-api-slim/hindsight_api/config.py +++ b/hindsight-api-slim/hindsight_api/config.py @@ -2072,7 +2072,17 @@ class LLMMemberConfig: # Valid multi-LLM strategy modes. LLM_STRATEGY_FAILOVER = "failover" LLM_STRATEGY_ROUND_ROBIN = "round-robin" -_VALID_LLM_STRATEGY_MODES = (LLM_STRATEGY_FAILOVER, LLM_STRATEGY_ROUND_ROBIN) +LLM_STRATEGY_METADATA = "metadata" +_VALID_LLM_STRATEGY_MODES = (LLM_STRATEGY_FAILOVER, LLM_STRATEGY_ROUND_ROBIN, LLM_STRATEGY_METADATA) + + +@dataclass(frozen=True) +class LLMMetadataRoute: + """One metadata match that pins an LLM request to a member.""" + + key: str + value: str + member: int @dataclass @@ -2080,22 +2090,27 @@ class LLMStrategyConfig: """How to route a request across the members of a multi-LLM chain. ``mode`` is "failover" (try members in order) or "round-robin" (rotate the - starting member per request, then fall through the rest on error). ``weights`` - is round-robin only: positive integers, one per member (primary first), giving - an unbalanced rotation; ``None`` means uniform. + starting member per request, then fall through the rest on error), or + "metadata" (pin when all matching routes agree, defaulting to the primary). + ``weights`` is round-robin only: positive integers, one per member (primary + first), giving an unbalanced rotation; ``routes`` is metadata only and every + matching route must select the same member. """ mode: str weights: list[int] | None = None + routes: list[LLMMetadataRoute] | None = None def _parse_llm_strategy(raw: str | None) -> LLMStrategyConfig | None: """Parse a multi-LLM strategy from a JSON env var. Returns ``None`` when unset. The value must be a JSON object with a ``mode`` - of "failover" or "round-robin"; ``weights`` (round-robin only) must be a list - of positive ints. Raises ``ValueError`` on any malformed input so - misconfiguration fails fast at startup rather than silently degrading. + of "failover", "round-robin", or "metadata". ``weights`` (round-robin only) + must be a list of positive ints. ``routes`` (metadata only) is a list + of ``{"key": str, "value": str, "member": non-negative int}`` objects. + Raises ``ValueError`` on malformed input so misconfiguration fails fast at + startup rather than silently degrading. """ text = (raw or "").strip() if not text: @@ -2118,7 +2133,29 @@ def _parse_llm_strategy(raw: str | None) -> LLMStrategyConfig | None: if not isinstance(weights, list) or not weights or not all(isinstance(w, int) and w > 0 for w in weights): raise ValueError("LLM strategy 'weights' must be a non-empty list of positive integers.") - return LLMStrategyConfig(mode=mode, weights=weights) + raw_routes = parsed.get("routes") + routes: list[LLMMetadataRoute] | None = None + if mode == LLM_STRATEGY_METADATA: + if not isinstance(raw_routes, list) or not raw_routes: + raise ValueError("LLM strategy 'routes' must be a non-empty list in metadata mode.") + routes = [] + for index, route in enumerate(raw_routes): + if not isinstance(route, dict): + raise ValueError(f"LLM metadata route {index} must be a JSON object.") + key = route.get("key") + value = route.get("value") + member = route.get("member") + if not isinstance(key, str) or not key: + raise ValueError(f"LLM metadata route {index} 'key' must be a non-empty string.") + if not isinstance(value, str): + raise ValueError(f"LLM metadata route {index} 'value' must be a string.") + if not isinstance(member, int) or isinstance(member, bool) or member < 0: + raise ValueError(f"LLM metadata route {index} 'member' must be a non-negative integer.") + routes.append(LLMMetadataRoute(key=key, value=value, member=member)) + elif raw_routes is not None: + raise ValueError(f"LLM strategy 'routes' is only valid with mode '{LLM_STRATEGY_METADATA}'.") + + return LLMStrategyConfig(mode=mode, weights=weights, routes=routes) def _parse_llm_members(prefix: str) -> list[LLMMemberConfig]: diff --git a/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py b/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py index e100953e7e..2132929852 100644 --- a/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py +++ b/hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py @@ -1342,7 +1342,13 @@ async def run_consolidation_job( # Build a configured LLM wrapper that applies per-bank settings (e.g. safety settings) # to every call without leaking across operations. - llm_config = memory_engine._consolidation_llm_config.with_config(config, bank_id=bank_id, operation="consolidation") + routing_metadata = memory_engine._pending_consolidation_llm_routing_metadata() + llm_config = memory_engine._consolidation_llm_config.with_config( + config, + bank_id=bank_id, + operation="consolidation", + routing_metadata=routing_metadata, + ) # Bind the operation trace context for the whole run so the create/update DB # sites (deep inside _process_memory_batch) can accumulate the observations @@ -2266,8 +2272,14 @@ async def _process_memory_batch( # "consolidation_dedup" (routes through the consolidation concurrency bucket via llm_wrapper's # "consolidation" prefix; recorded distinctly in llm_requests). dedup_enabled = _dedup_active(config) + routing_metadata = memory_engine._pending_consolidation_llm_routing_metadata() if dedup_enabled else None dedup_llm_config = ( - memory_engine._consolidation_llm_config.with_config(config, bank_id=bank_id, operation="consolidation_dedup") + memory_engine._consolidation_llm_config.with_config( + config, + bank_id=bank_id, + operation="consolidation_dedup", + routing_metadata=routing_metadata, + ) if dedup_enabled else None ) diff --git a/hindsight-api-slim/hindsight_api/engine/llm_wrapper.py b/hindsight-api-slim/hindsight_api/engine/llm_wrapper.py index e9912d1022..b4bb6052e8 100644 --- a/hindsight-api-slim/hindsight_api/engine/llm_wrapper.py +++ b/hindsight-api-slim/hindsight_api/engine/llm_wrapper.py @@ -1595,6 +1595,7 @@ def with_config( bank_id: str | None = None, operation: str | None = None, metadata: dict[str, Any] | None = None, + routing_metadata: dict[str, Any] | None = None, ) -> "ConfiguredLLMProvider": """ Return a configured wrapper for a specific bank operation. @@ -1609,6 +1610,8 @@ def with_config( operation: Logical operation label ("retain", "reflect", ...) for LLM trace rows. metadata: Optional extra caller metadata stored on trace rows. + routing_metadata: Optional ephemeral values used by a multi-LLM + router. Ignored by a single provider and never persisted. Returns: A ``ConfiguredLLMProvider`` that delegates to this provider with diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py index b1c54991d3..76645aa409 100644 --- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py +++ b/hindsight-api-slim/hindsight_api/engine/memory_engine.py @@ -47,6 +47,7 @@ DEFAULT_REFLECT_SOURCE_FACTS_MAX_TOKENS, DEFAULT_STORE_DOCUMENT_TEXT, ENV_MODEL_INIT_TIMEOUT, + LLM_STRATEGY_METADATA, HindsightConfig, LLMMemberConfig, LLMStrategyConfig, @@ -683,6 +684,48 @@ def _build_llm( return MultiLLMProvider([base, *extra], strategy) +@dataclass(frozen=True) +class _AppendRoutingState: + """Stored classification that an append must route and optionally inherit.""" + + tags: list[str] + metadata: dict[str, Any] + + +def _retain_llm_routing_metadata( + contents: list[RetainContentDict], + document_tags: list[str] | None, + append_states: list[_AppendRoutingState] | None = None, +) -> dict[str, list[str]]: + """Build ephemeral metadata used to select a retain LLM member. + + A retain LLM prompt may combine several items, so each key carries the + union of values present in the batch. Matching any sensitive item therefore + routes the whole prompt to the sensitive lane instead of exposing its peers + to the primary. User metadata is namespaced to keep the reserved ``tags`` + route key unambiguous. + """ + values: dict[str, set[str]] = {} + tags = set(document_tags or []) + for state in append_states or []: + tags.update(state.tags) + for key, value in state.metadata.items(): + route_values = value if isinstance(value, (list, tuple, set, frozenset)) else [value] + for route_value in route_values: + if route_value is not None: + values.setdefault(f"metadata.{key}", set()).add(str(route_value)) + for item in contents: + tags.update(item.get("tags", []) or []) + for key, value in (item.get("metadata") or {}).items(): + route_values = value if isinstance(value, (list, tuple, set, frozenset)) else [value] + for route_value in route_values: + if route_value is not None: + values.setdefault(f"metadata.{key}", set()).add(str(route_value)) + if tags: + values["tags"] = tags + return {key: sorted(route_values) for key, route_values in values.items()} + + async def validate_retain_batch_support( retain_llm_config: "LLMConfig | MultiLLMProvider", config: HindsightConfig ) -> None: @@ -690,9 +733,9 @@ async def validate_retain_batch_support( Otherwise the server would silently fall back to sync mode on every retain, which is confusing and wastes a config knob. For a multi-LLM chain the - capability is evaluated across ALL members, not just the primary: batch - capacity may live on a secondary (issue #3645), and gating on the primary - alone rejected configurations that would in fact have worked. + failover/round-robin capacity is evaluated across all members because it may + live on a secondary (issue #3645). Metadata mode instead requires every + route-selectable member because each request is pinned before submission. """ if not config.retain_batch_enabled: return @@ -700,8 +743,15 @@ async def validate_retain_batch_support( return if isinstance(retain_llm_config, MultiLLMProvider): - members = ", ".join(f"'{member.provider}'" for member in retain_llm_config.members) - cause = f"no member of the retain LLM chain ({members}) supports the batch API" + if retain_llm_config.strategy.mode == LLM_STRATEGY_METADATA: + selectable_indices = {0, *(route.member for route in retain_llm_config.strategy.routes or [])} + members = ", ".join( + f"'{retain_llm_config.members[index].provider}'" for index in sorted(selectable_indices) + ) + cause = f"metadata routing requires every selectable retain LLM member ({members}) to support the batch API" + else: + members = ", ".join(f"'{member.provider}'" for member in retain_llm_config.members) + cause = f"no member of the retain LLM chain ({members}) supports the batch API" else: cause = f"the retain LLM provider '{retain_llm_config.provider}' does not support the batch API" raise RuntimeError( @@ -2861,7 +2911,7 @@ async def _handle_file_convert_retain(self, task_dict: dict[str, Any]): document_tags = task_dict.get("document_tags") retain_task_payload: dict[str, Any] = {"contents": retain_contents} - if document_tags: + if document_tags is not None: retain_task_payload["document_tags"] = document_tags if task_dict.get("strategy"): retain_task_payload["strategy"] = task_dict["strategy"] @@ -5146,6 +5196,45 @@ async def retain_batch_async( "document text is not stored and cannot be appended to. Use update_mode='replace' instead." ) + # Append re-extracts the stored document body along with the new text. + # Capture its classification once, before the first sub-batch can + # overwrite the document record, so every slice stays on the same + # metadata-routed LLM lane even when the append omits those fields. + stored_append_states = await self._stored_append_routing_states(bank_id, contents) + + # Append updates inherit stored classification keys so adding unrelated + # tags/metadata cannot silently erase the routing policy for the next + # append. An explicitly supplied empty collection remains an intentional + # declassification; the current append still routes with the old state + # because its prompt reprocesses the old body, while later appends see + # the cleared state. + for item in contents: + if item.get("update_mode") != "append" or not item.get("document_id"): + continue + stored_state = stored_append_states.get(item["document_id"]) + if stored_state is None: + continue + item_tags = item.get("tags") + tags_explicitly_cleared = item_tags == [] or document_tags == [] + if stored_state.tags and not tags_explicitly_cleared: + item["tags"] = list(dict.fromkeys([*stored_state.tags, *(item_tags or [])])) + + item_metadata = item.get("metadata") + if stored_state.metadata and item_metadata != {}: + item["metadata"] = {**stored_state.metadata, **(item_metadata or {})} + + # Validate the complete submission before token splitting or per-document + # grouping can hide a cross-member classification conflict. No LLM call + # may start when one logical retain operation belongs to multiple lanes. + if isinstance(self._retain_llm_config, MultiLLMProvider): + self._retain_llm_config.validate_routing_metadata( + _retain_llm_routing_metadata( + contents, + document_tags, + list(stored_append_states.values()), + ) + ) + # Fold items that share an explicit document_id into one document. On the # synchronous in-process path this is safe — sub-batches run sequentially, # so same-document items cannot race (unlike the queued path, which still @@ -5171,6 +5260,7 @@ async def retain_batch_async( document_id=document_id, fact_type_override=fact_type_override, document_tags=document_tags, + routing_append_states=list(stored_append_states.values()), operation_id=operation_id, strategy=strategy, outbox_callback=outbox_callback, @@ -5233,6 +5323,11 @@ async def retain_batch_async( is_first_batch=True, fact_type_override=fact_type_override, document_tags=document_tags, + routing_append_states=( + [stored_append_states[group.document_id]] + if group.document_id is not None and group.document_id in stored_append_states + else [] + ), operation_id=operation_id, strategy=strategy, outbox_callback=group_outbox_callback, @@ -5460,6 +5555,113 @@ def _retain_chunking_config(config: HindsightConfig) -> _RetainChunkingConfig: structured_chunk_size=config.retain_structured_chunk_size, ) + async def _stored_append_routing_states( + self, + bank_id: str, + contents: list[RetainContentDict], + ) -> dict[str, _AppendRoutingState]: + """Read existing append classification before retain can replace it. + + This is only needed for metadata-routed retain chains. Other strategies + keep their existing append path and do not pay for an extra database read. + """ + if not ( + isinstance(self._retain_llm_config, MultiLLMProvider) + and self._retain_llm_config.strategy.mode == LLM_STRATEGY_METADATA + ): + return {} + + append_document_ids = { + item["document_id"] for item in contents if item.get("update_mode") == "append" and item.get("document_id") + } + if not append_document_ids: + return {} + + document_ids = sorted(append_document_ids) + stored_states: dict[str, _AppendRoutingState] = {} + backend = await self._get_backend() + from .memories import get_memories + + store = get_memories() + async with acquire_with_retry(backend) as conn: + rows = await conn.fetch( + f"SELECT id, tags, retain_params FROM {fq_table('documents')} " + f"WHERE id = ANY($1::text[]) AND bank_id = $2", + document_ids, + bank_id, + ) + for row in rows: + tags = conn.parse_json(row["tags"]) + retain_params = conn.parse_json(row["retain_params"]) + if tags is None and retain_params is None: + continue + metadata = retain_params.get("metadata", {}) if isinstance(retain_params, dict) else {} + stored_states[str(row["id"])] = _AppendRoutingState( + tags=list(tags or []), + metadata=dict(metadata) if isinstance(metadata, dict) else {}, + ) + + # Release the pooled database connection before crossing the external + # store seam. Store-owned records are authoritative and may not have a + # SQL mirror; use the bulk API so remote stores can serve all documents + # in one round trip. + if store.store_owned_for(bank_id): + records = await store.get_document_records(bank_id=bank_id, document_ids=document_ids) + for append_document_id, record in records.items(): + existing = stored_states.get(append_document_id) + tags = record.get("tags") + if tags is None and existing is not None: + tags = existing.tags + record_metadata = record.get("metadata") or {} + retain_params = record_metadata.get("retain_params") + if retain_params is None and existing is not None: + metadata = existing.metadata + else: + retain_params = json.loads(retain_params) if isinstance(retain_params, str) else retain_params + metadata = retain_params.get("metadata", {}) if isinstance(retain_params, dict) else {} + if tags is None and not metadata: + continue + stored_states[append_document_id] = _AppendRoutingState( + tags=list(tags or []), + metadata=dict(metadata) if isinstance(metadata, dict) else {}, + ) + return stored_states + + def _reflect_llm_routing_metadata( + self, + llm_config: "LLMConfig | MultiLLMProvider", + ) -> dict[str, list[str]]: + """Bind reflect to the configured protected tag lane. + + Reflect selects its LLM before the agent runs retrieval tools. Even an + exact tag scope can expand an allowed fact to its full document, whose + other facts may carry a protected tag. A preflight existence query would + also race with concurrent retain. Include every configured tag route so + the whole agent loop and mental-model delta calls stay fail-closed. + + Non-metadata strategies return without additional work. + """ + if not (isinstance(llm_config, MultiLLMProvider) and llm_config.strategy.mode == LLM_STRATEGY_METADATA): + return {} + + route_tags = {route.value for route in llm_config.strategy.routes or [] if route.key == "tags"} + return {"tags": sorted(route_tags)} if route_tags else {} + + def _pending_consolidation_llm_routing_metadata(self) -> dict[str, list[str]]: + """Bind consolidation fail-closed against facts retained mid-run. + + A consolidation job repeatedly fetches pending batches. A sensitive fact + may arrive after startup, so a point-in-time existence check cannot safely + select the primary for the whole job. Including every configured tag route + keeps all batches on the configured protected lane. + """ + llm_config = self._consolidation_llm_config + if not (isinstance(llm_config, MultiLLMProvider) and llm_config.strategy.mode == LLM_STRATEGY_METADATA): + return {} + + route_tags = {route.value for route in llm_config.strategy.routes or [] if route.key == "tags"} + return {"tags": sorted(route_tags)} if route_tags else {} + async def _run_retain_execution( self, *, @@ -5469,6 +5671,7 @@ async def _run_retain_execution( document_id: str | None, fact_type_override: str | None, document_tags: list[str] | None, + routing_append_states: list[_AppendRoutingState], operation_id: str | None, strategy: str | None, outbox_callback: RetainOutboxCallback | None, @@ -5654,6 +5857,7 @@ async def _run_sub(idx: int, contents_, origins_, offset_, is_last_, body_, body is_first_batch=idx == 1, # Only upsert on first batch fact_type_override=fact_type_override, document_tags=document_tags, + routing_append_states=routing_append_states, operation_id=operation_id, strategy=strategy, # Outbox callback runs inside the last sub-batch's transaction so the @@ -5828,6 +6032,7 @@ async def _run_sub(idx: int, contents_, origins_, offset_, is_last_, body_, body is_first_batch=True, fact_type_override=fact_type_override, document_tags=document_tags, + routing_append_states=routing_append_states, operation_id=operation_id, strategy=strategy, outbox_callback=outbox_callback, @@ -5877,6 +6082,7 @@ async def _retain_batch_async_internal( is_first_batch: bool = True, fact_type_override: str | None = None, document_tags: list[str] | None = None, + routing_append_states: list[_AppendRoutingState] | None = None, operation_id: str | None = None, outbox_callback: RetainOutboxCallback | None = None, outbox_callback_factory: RetainOutboxCallbackFactory | None = None, @@ -5903,6 +6109,8 @@ async def _retain_batch_async_internal( is_first_batch: Whether this is the first batch (for chunked operations, only delete on first batch) fact_type_override: Override fact type for all facts document_tags: Tags applied to all items in this batch + routing_append_states: Existing append-document classification used + for LLM selection; omitted fields also inherit it before this call Returns: Tuple of (unit ID lists, LLM token usage, processed_content_tokens). @@ -5928,7 +6136,16 @@ async def _retain_batch_async_internal( # Create parent span for retain operation with create_operation_span("retain", bank_id): - retain_llm = self._retain_llm_config.with_config(resolved_config, bank_id=bank_id, operation="retain") + retain_llm = self._retain_llm_config.with_config( + resolved_config, + bank_id=bank_id, + operation="retain", + routing_metadata=_retain_llm_routing_metadata( + contents, + document_tags, + routing_append_states, + ), + ) result = await self._retain_batch_with_append_retry( pool=self._backend, embeddings_model=self.embeddings, @@ -13023,6 +13240,7 @@ async def reflect_async( # This is critical for banks with many mental models to avoid huge prompts. resolved_reflect_config = await self._config_resolver.resolve_full_config(bank_id, request_context) + reflect_routing_metadata = self._reflect_llm_routing_metadata(self._reflect_llm_config) # Compute max iterations based on budget config = get_config() @@ -13207,7 +13425,10 @@ async def expand_fn(memory_ids: list[str], depth: str) -> dict[str, Any]: agent_result = await asyncio.wait_for( run_reflect_agent( llm_config=self._reflect_llm_config.with_config( - resolved_reflect_config, bank_id=bank_id, operation=_operation_label + resolved_reflect_config, + bank_id=bank_id, + operation=_operation_label, + routing_metadata=reflect_routing_metadata, ), bank_id=bank_id, query=query, @@ -15344,11 +15565,13 @@ async def _op_llm() -> ConfiguredLLMProvider: nonlocal _op_llm_config if _op_llm_config is None: resolved_config = await self._config_resolver.resolve_full_config(bank_id, request_context) + routing_metadata = self._reflect_llm_routing_metadata(self._reflect_llm_config) _op_llm_config = self._reflect_llm_config.with_config( resolved_config, bank_id=bank_id, operation="mental_model_delta_ops", metadata={"mental_model_id": str(mental_model_id)}, + routing_metadata=routing_metadata, ) return _op_llm_config @@ -19504,6 +19727,23 @@ async def submit_async_retain( f"batch synchronously (async=false), which processes them sequentially." ) + # Validate the parent payload before child packing, database inserts, or + # task dispatch. Otherwise classifications routed to different members + # can land in separate workers and each appear unambiguous in isolation. + routing_contents = cast(list[RetainContentDict], contents) + if ( + isinstance(self._retain_llm_config, MultiLLMProvider) + and self._retain_llm_config.strategy.mode == LLM_STRATEGY_METADATA + ): + stored_append_states = await self._stored_append_routing_states(bank_id, routing_contents) + self._retain_llm_config.validate_routing_metadata( + _retain_llm_routing_metadata( + routing_contents, + document_tags, + list(stored_append_states.values()), + ) + ) + # Calculate total token count and determine if we need to split total_tokens = sum(count_tokens(item.get("content", "")) for item in contents) config = get_config() @@ -19516,7 +19756,7 @@ async def submit_async_retain( # on the same document_id and trigger FK violations in the final # ANN pass (issue #1795). The worker's in-process splitter # handles intra-document chunking sequentially. - sub_batches = _split_contents_into_async_children(cast(list[RetainContentDict], contents), tokens_per_batch) + sub_batches = _split_contents_into_async_children(routing_contents, tokens_per_batch) # Log splitting info if we actually split if len(sub_batches) > 1: @@ -19606,7 +19846,7 @@ async def submit_async_retain( ) task_payload: dict[str, Any] = {"contents": sub_batch} - if document_tags: + if document_tags is not None: task_payload["document_tags"] = document_tags if strategy: task_payload["strategy"] = strategy diff --git a/hindsight-api-slim/hindsight_api/engine/multi_llm.py b/hindsight-api-slim/hindsight_api/engine/multi_llm.py index 7163504a87..24256330e8 100644 --- a/hindsight-api-slim/hindsight_api/engine/multi_llm.py +++ b/hindsight-api-slim/hindsight_api/engine/multi_llm.py @@ -1,4 +1,4 @@ -"""Multi-LLM routing: failover and (weighted) round-robin across N providers. +"""Multi-LLM routing across N providers. ``MultiLLMProvider`` wraps an ordered list of :class:`LLMProvider` members and a :class:`~hindsight_api.config.LLMStrategyConfig`, exposing the same public surface @@ -14,13 +14,16 @@ - ``failover``: try members in declared order ``[0..N]``. - ``round-robin``: rotate the starting member per request (optionally weighted), then fall through the remaining members on error. - -Batch retain runs on the **first batch-capable member** in declared order (see -``batch_provider_impl``), which need not be the primary; once selected, the whole -batch lifecycle stays on that member and does not fail over. Every other direct -``_provider_impl`` access still resolves to the primary via attribute passthrough -— failover/round-robin apply to the interactive ``call`` / ``call_with_tools`` -paths. +- ``metadata``: bind an operation to the member selected by its ephemeral + metadata, defaulting to member 0. Matches to different members are rejected; + a selected member is strict and never falls across lanes on error. + +For failover/round-robin, batch retain runs on the **first batch-capable member** +in declared order (see ``batch_provider_impl``), which need not be the primary. +Metadata routing pins the member before batch submission. Once selected, the +whole batch lifecycle stays on that member and does not fail over. Every other +direct ``_provider_impl`` access still resolves to the primary via attribute +passthrough. """ import logging @@ -28,7 +31,7 @@ import uuid from typing import TYPE_CHECKING, Any -from ..config import LLM_STRATEGY_FAILOVER, LLMStrategyConfig +from ..config import LLM_STRATEGY_FAILOVER, LLM_STRATEGY_METADATA, LLMStrategyConfig from .llm_wrapper import LLMProvider, OutputTooLongError if TYPE_CHECKING: @@ -97,16 +100,65 @@ def __init__(self, members: list[LLMProvider], strategy: LLMStrategyConfig) -> N ) self._scheduler = _WeightedRoundRobin(weights) + if strategy.mode == LLM_STRATEGY_METADATA: + if not strategy.routes: + raise ValueError("Metadata LLM routing requires at least one route") + for route in strategy.routes: + if route.member >= len(members): + raise ValueError( + f"LLM metadata route {route.key}={route.value!r} selects member {route.member}, " + f"but the chain has members 0..{len(members) - 1}." + ) + tag_members = {route.member for route in strategy.routes if route.key == "tags"} + if len(tag_members) > 1: + raise ValueError( + "Metadata LLM routes with key 'tags' must all select the same member; " + "reflect and consolidation include every configured tag route to stay fail-closed." + ) + # ── routing ──────────────────────────────────────────────────────────────── def _member_order(self) -> list[int]: """Indices to try, in order, for one request.""" n = len(self._members) + if self._strategy.mode == LLM_STRATEGY_METADATA: + # Metadata is supplied at operation binding time (``with_config``). + # A direct call has no routing context, so it stays on the primary. + return [0] if self._strategy.mode == LLM_STRATEGY_FAILOVER: return list(range(n)) start = self._scheduler.next() return [(start + i) % n for i in range(n)] + def _member_for_metadata(self, metadata: dict[str, Any] | None) -> LLMProvider: + """Return the unambiguous metadata route match, or the primary. + + One LLM call cannot be split across classification lanes. If its combined + metadata matches routes to different members, choosing either member + could disclose the other lane's data. Refuse the operation instead. + """ + matched_members: list[int] = [] + for route in self._strategy.routes or []: + actual = (metadata or {}).get(route.key) + matches = actual == route.value or ( + isinstance(actual, (list, tuple, set, frozenset)) and route.value in actual + ) + if matches: + matched_members.append(route.member) + + distinct_members = set(matched_members) + if len(distinct_members) > 1: + raise ValueError( + "LLM metadata routes for this operation select multiple members; " + "split the input or route all matching classifications to one member." + ) + return self._members[matched_members[0]] if matched_members else self._members[0] + + def validate_routing_metadata(self, metadata: dict[str, Any] | None) -> None: + """Validate operation-wide metadata before work is split or queued.""" + if self._strategy.mode == LLM_STRATEGY_METADATA: + self._member_for_metadata(metadata) + async def _dispatch(self, method_name: str, **kwargs: Any) -> Any: last_exc: BaseException | None = None order = self._member_order() @@ -148,10 +200,10 @@ async def call_with_tools( async def verify_connection(self) -> None: """Strictly verify the primary; soft-verify the rest (warn, don't fail). - A failover member being unreachable at startup must not block the server — - it may come back before it's needed. The primary is the steady-state path, - so its failure is still surfaced (the caller already wraps this in a - warn-only try/except at startup). + A secondary being unreachable at startup must not block the server — it + may come back before its routing strategy selects it. The primary is the + steady-state path, so its failure is still surfaced (the caller already + wraps this in a warn-only try/except at startup). """ await self._members[0].verify_connection() for member in self._members[1:]: @@ -159,8 +211,8 @@ async def verify_connection(self) -> None: await member.verify_connection() except Exception as e: # noqa: BLE001 - soft verification logger.warning( - "Failover LLM member %s/%s failed connection verification: %s. " - "It will be tried at request time if the primary fails.", + "Secondary LLM member %s/%s failed connection verification: %s. " + "It will be used if the routing strategy selects it at request time.", member.provider, member.model, e, @@ -169,13 +221,18 @@ async def verify_connection(self) -> None: # ── batch routing ─────────────────────────────────────────────────────────── async def supports_batch_api(self) -> bool: - """Whether ANY member supports the batch API. + """Whether the configured strategy can safely use the batch API. - The single-provider path delegates to the primary, but in a multi-LLM - chain batch capacity may live on a secondary member (e.g. an ``openai`` / - ``groq`` fallback behind a non-batch primary). Mirroring the failover - semantics, the batch path can proceed as long as one member can serve it. + Failover and round-robin need any one capable member. Metadata routing + needs every selectable member because the request-specific member is not + known during startup validation. """ + if self._strategy.mode == LLM_STRATEGY_METADATA: + routed_indices = {0, *(route.member for route in self._strategy.routes or [])} + for index in routed_indices: + if not await self._members[index].supports_batch_api(): + return False + return True return (await self.batch_provider_impl()) is not None async def batch_provider_impl(self, account_key: str | None = None) -> "LLMInterface | None": @@ -195,7 +252,8 @@ async def batch_provider_impl(self, account_key: str | None = None) -> "LLMInter submit time and gets back the member that owns the batch, or ``None`` — never a lookalike. """ - for member in self._members: + members = self._members[:1] if self._strategy.mode == LLM_STRATEGY_METADATA else self._members + for member in members: impl = await member.batch_provider_impl(account_key) if impl is not None: return impl @@ -212,10 +270,22 @@ def with_config( bank_id: str | None = None, operation: str | None = None, metadata: dict[str, Any] | None = None, + routing_metadata: dict[str, Any] | None = None, ) -> "ConfiguredLLMProvider": """Mirror ``LLMProvider.with_config`` so the strategy runs inside the per-operation configured wrapper (gemini-safety + trace contextvars wrap every member call).""" + if self._strategy.mode == LLM_STRATEGY_METADATA: + # Pin the whole operation before any interactive or batch call. A + # sensitive route is intentionally strict: if its selected provider + # fails, its content must not spill into a differently classified lane. + return self._member_for_metadata(routing_metadata).with_config( + config, + bank_id=bank_id, + operation=operation, + metadata=metadata, + ) + from .llm_trace import LLMTraceContext from .llm_wrapper import ConfiguredLLMProvider @@ -236,6 +306,10 @@ def with_config( def members(self) -> list[LLMProvider]: return self._members + @property + def strategy(self) -> LLMStrategyConfig: + return self._strategy + def __getattr__(self, name: str) -> Any: # Anything not defined here (provider, model, api_key, base_url, # _provider_impl, mock helpers, ...) delegates to the primary member so diff --git a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py index cf52bcf626..bc8251ea9e 100644 --- a/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py +++ b/hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py @@ -377,7 +377,7 @@ def _build_retain_params(contents_dicts, document_tags=None, doc_contents=None): if items: first_item = items[0] for key, value in first_item.items(): - if key in _RETAIN_PARAMS_NOT_REPLAYED or value is None: + if key in _RETAIN_PARAMS_NOT_REPLAYED or key == "metadata" or value is None: continue # event_date arrives as a datetime from some callers and a string from # others; retain_params is JSON, so normalise here rather than at every @@ -386,6 +386,28 @@ def _build_retain_params(contents_dicts, document_tags=None, doc_contents=None): value = value.isoformat() if hasattr(value, "isoformat") else str(value) retain_params[key] = value + metadata_values: dict[str, list[Any]] = {} + for item in items: + item_metadata = item.get("metadata") + if not isinstance(item_metadata, dict): + continue + for key, value in item_metadata.items(): + values = metadata_values.setdefault(key, []) + if value not in values: + values.append(value) + if metadata_values: + # A shared document can combine several input items. Persist every + # classification value so an append cannot inherit only items[0] + # and accidentally fall back to the primary LLM. Single-valued + # keys retain their existing scalar representation. + retain_params["metadata"] = { + key: values[0] if len(values) == 1 else values for key, values in metadata_values.items() + } + elif first_item.get("metadata") is not None: + # Validation normally guarantees a mapping. Keep the historical + # pass-through for internal callers and the round-trip contract. + retain_params["metadata"] = first_item["metadata"] + return retain_params, merged_tags diff --git a/hindsight-api-slim/tests/test_async_retain_tags.py b/hindsight-api-slim/tests/test_async_retain_tags.py index 45ef15ccc1..f2af5a75b5 100644 --- a/hindsight-api-slim/tests/test_async_retain_tags.py +++ b/hindsight-api-slim/tests/test_async_retain_tags.py @@ -9,7 +9,8 @@ @pytest.mark.asyncio -async def test_submit_async_retain_includes_document_tags_in_task_payload(): +@pytest.mark.parametrize("document_tags", [["scope:tools", "user:alice"], []]) +async def test_submit_async_retain_includes_document_tags_in_task_payload(document_tags): """submit_async_retain should include document_tags in queued task payload. submit_async_batch_retain inserts the parent + all children inline inside @@ -25,6 +26,7 @@ async def test_submit_async_retain_includes_document_tags_in_task_payload(): engine._initialized = True engine._authenticate_tenant = AsyncMock() engine._operation_validator = None + engine._retain_llm_config = MagicMock() # Children are now inserted inline (no _submit_async_operation hop), and # submit_task fires post-commit. Mock both so the inline path runs cleanly # without the test needing real DB or task backend. @@ -54,7 +56,6 @@ async def test_submit_async_retain_includes_document_tags_in_task_payload(): request_context = RequestContext(tenant_id="tenant-a", api_key_id="key-a") contents = [{"content": "Async retain payload test."}] - document_tags = ["scope:tools", "user:alice"] # Stub the lazy bank-create/default-template hook to a no-op (created=False) # so the inline transaction path runs against the mock connection without diff --git a/hindsight-api-slim/tests/test_consolidation_dedup.py b/hindsight-api-slim/tests/test_consolidation_dedup.py index 6f6ba66641..c08ca3dd18 100644 --- a/hindsight-api-slim/tests/test_consolidation_dedup.py +++ b/hindsight-api-slim/tests/test_consolidation_dedup.py @@ -11,7 +11,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime, timezone -from unittest.mock import DEFAULT, AsyncMock, patch +from unittest.mock import DEFAULT, AsyncMock, MagicMock, patch import pytest @@ -668,17 +668,34 @@ async def test_dedup_update_all_updated_sources_deleted_skips_fold_and_delete() # ── _process_memory_batch create-contract (created vs skipped) ──────────────── -def _batch_engine(): - return types.SimpleNamespace(_consolidation_llm_config=types.SimpleNamespace(with_config=lambda *a, **k: object())) - - -async def _run_create_batch(create_action_result: str): +@dataclass +class _CreateBatchRun: + result: tuple[list[dict], int, bool] + create_action: AsyncMock + memory_id: str + base_llm_config: MagicMock + child_llm_config: object + config: object + routing_metadata: dict[str, list[str]] + dedup_adjudicate: AsyncMock + + +async def _run_create_batch(create_action_result: str) -> _CreateBatchRun: from hindsight_api.engine.consolidation import consolidator as C mem_id = str(uuid.uuid4()) memories = [{"id": mem_id, "text": "Uzbek YouTube content is very rich.", "tags": []}] create = C._CreateAction(text="Uzbek YouTube content is very rich.", source_fact_ids=[mem_id]) llm_result = C._BatchLLMResult(creates=[create]) + child_llm_config = object() + base_llm_config = MagicMock() + base_llm_config.with_config.return_value = child_llm_config + routing_metadata = {"tags": ["sensitive"]} + memory_engine = types.SimpleNamespace( + _consolidation_llm_config=base_llm_config, + _pending_consolidation_llm_routing_metadata=lambda: routing_metadata, + ) + config = object() with ( patch.object( C, @@ -696,40 +713,62 @@ async def _run_create_batch(create_action_result: str): C, "_dedup_adjudicate", new=AsyncMock(return_value=_DedupOutcome(best_id=None, merged_text="", should_merge=False)), - ), + ) as dedup_adjudicate, patch.object(C, "_apply_create_action", new=AsyncMock(return_value=create_action_result)) as create_action, ): result = await C._process_memory_batch( pool=_DedupBackend(_DedupConn()), - memory_engine=_batch_engine(), + memory_engine=memory_engine, llm_config=object(), bank_id="bank1", memories=memories, request_context=object(), - config=object(), + config=config, ) - return result, create_action, mem_id + return _CreateBatchRun( + result=result, + create_action=create_action, + memory_id=mem_id, + base_llm_config=base_llm_config, + child_llm_config=child_llm_config, + config=config, + routing_metadata=routing_metadata, + dedup_adjudicate=dedup_adjudicate, + ) async def test_process_batch_creates_when_dedup_target_vanished() -> None: # Caller contract: when the adjudicator finds no twin to fold into, _process_memory_batch # must still CREATE the observation instead of dropping it. - result, create_action, mem_id = await _run_create_batch("created") - create_action.assert_awaited_once() - prepared = create_action.await_args.kwargs["prepared"] + run = await _run_create_batch("created") + run.create_action.assert_awaited_once() + prepared = run.create_action.await_args.kwargs["prepared"] assert prepared.text == "Uzbek YouTube content is very rich." - assert prepared.source_memory_ids == [mem_id] - assert result == ([{"action": "created"}], 0, False) + assert prepared.source_memory_ids == [run.memory_id] + assert run.result == ([{"action": "created"}], 0, False) async def test_process_batch_reports_skipped_when_create_skipped() -> None: # _apply_create_action returns "skipped" (all sources deleted in the write txn) -> # _process_memory_batch must NOT mark the memory created; it falls through to skipped. - result, _create_action, _mem_id = await _run_create_batch("skipped") - assert result == ([{"action": "skipped", "reason": "no_durable_knowledge"}], 0, False) + run = await _run_create_batch("skipped") + assert run.result == ([{"action": "skipped", "reason": "no_durable_knowledge"}], 0, False) async def test_process_batch_reports_created_when_create_created() -> None: # _apply_create_action returns "created" -> the memory is marked created. - result, _create_action, _mem_id = await _run_create_batch("created") - assert result == ([{"action": "created"}], 0, False) + run = await _run_create_batch("created") + assert run.result == ([{"action": "created"}], 0, False) + + +async def test_process_batch_derives_dedup_wrapper_from_pinned_consolidation_lane() -> None: + run = await _run_create_batch("created") + + assert run.result == ([{"action": "created"}], 0, False) + run.base_llm_config.with_config.assert_called_once_with( + run.config, + bank_id="bank1", + operation="consolidation_dedup", + routing_metadata=run.routing_metadata, + ) + assert run.dedup_adjudicate.await_args.args[4] is run.child_llm_config diff --git a/hindsight-api-slim/tests/test_file_retain.py b/hindsight-api-slim/tests/test_file_retain.py index 0f9b2dafa7..8461958a66 100644 --- a/hindsight-api-slim/tests/test_file_retain.py +++ b/hindsight-api-slim/tests/test_file_retain.py @@ -1013,7 +1013,8 @@ async def run_case(label: str, timestamp_value) -> dict: @pytest.mark.asyncio -async def test_file_retain_forwards_all_content_fields(memory_no_llm_verify, sample_txt_content): +@pytest.mark.parametrize("document_tags", [["batch_tag"], []]) +async def test_file_retain_forwards_all_content_fields(memory_no_llm_verify, sample_txt_content, document_tags): """Regression: _handle_file_convert_retain must forward every FileRetainMetadata field to the inner batch_retain task without renaming or dropping it. @@ -1080,7 +1081,7 @@ async def capturing_submit(task_dict): "strategy": "my_strategy", } ], - document_tags=["batch_tag"], + document_tags=document_tags, request_context=request_context, ) @@ -1090,7 +1091,7 @@ async def capturing_submit(task_dict): # Per-request fields (live on the outer task payload, not per-content). assert payload.get("strategy") == "my_strategy", "strategy must be forwarded at request level" - assert payload.get("document_tags") == ["batch_tag"], "document_tags must be forwarded at request level" + assert payload.get("document_tags") == document_tags, "document_tags must be forwarded at request level" # Per-content fields. assert len(payload["contents"]) == 1 diff --git a/hindsight-api-slim/tests/test_mental_model_delta.py b/hindsight-api-slim/tests/test_mental_model_delta.py index c8833cf944..fda57f3d6f 100644 --- a/hindsight-api-slim/tests/test_mental_model_delta.py +++ b/hindsight-api-slim/tests/test_mental_model_delta.py @@ -28,8 +28,10 @@ import pytest from hindsight_api import MemoryEngine, RequestContext -from hindsight_api.engine.llm_wrapper import LLMConfig +from hindsight_api.config import LLMMetadataRoute, LLMStrategyConfig +from hindsight_api.engine.llm_wrapper import LLMConfig, LLMProvider from hindsight_api.engine.maintenance import MaintenanceLoop +from hindsight_api.engine.multi_llm import MultiLLMProvider from hindsight_api.engine.response_models import ReflectResult from hindsight_api.engine.retain import embedding_utils from tests.conftest import stub_refresh_has_sources @@ -1053,8 +1055,27 @@ async def test_delta_call_is_traced_and_uses_decoupled_completion_cap( patch_reflect(memory, text="ignored — full mode candidate") captured: dict[str, Any] = {} + primary_calls = 0 + secondary_calls = 0 + + primary = LLMProvider(provider="mock", api_key="", base_url="", model="primary") + secondary = LLMProvider(provider="mock", api_key="", base_url="", model="secondary") + memory._reflect_llm_config = MultiLLMProvider( + [primary, secondary], + LLMStrategyConfig( + mode="metadata", + routes=[LLMMetadataRoute(key="tags", value="sensitive", member=1)], + ), + ) + + async def primary_call(*, messages, **kwargs): + nonlocal primary_calls + primary_calls += 1 + raise AssertionError("mental-model delta must not use the primary lane") async def capturing_call(*, messages, **kwargs): + nonlocal secondary_calls + secondary_calls += 1 ctx = current_trace_context() captured["max_completion_tokens"] = kwargs.get("max_completion_tokens") captured["scope"] = kwargs.get("scope") @@ -1064,7 +1085,8 @@ async def capturing_call(*, messages, **kwargs): return DeltaOperationList() # First (seeding) refresh — value captured here is overwritten by the second. - monkeypatch.setattr(memory._reflect_llm_config, "call", capturing_call) + monkeypatch.setattr(primary, "call", primary_call) + monkeypatch.setattr(secondary, "call", capturing_call) await memory.refresh_mental_model(bank_id=bank_id, mental_model_id=mm["id"], request_context=request_context) # Second refresh with a genuine new fact so the delta call actually fires. @@ -1085,6 +1107,8 @@ async def capturing_call(*, messages, **kwargs): assert captured["trace_operation"] == "mental_model_delta_ops" assert captured["trace_bank_id"] == bank_id assert captured["trace_metadata"] == {"mental_model_id": str(mm["id"])} + assert secondary_calls == 1 + assert primary_calls == 0 await memory.delete_bank(bank_id, request_context=request_context) diff --git a/hindsight-api-slim/tests/test_metadata_llm_routing.py b/hindsight-api-slim/tests/test_metadata_llm_routing.py new file mode 100644 index 0000000000..d28cdc5f82 --- /dev/null +++ b/hindsight-api-slim/tests/test_metadata_llm_routing.py @@ -0,0 +1,641 @@ +"""End-to-end retain coverage for metadata-based LLM routing.""" + +import copy +import dataclasses +import uuid +from contextlib import asynccontextmanager +from unittest.mock import AsyncMock + +import pytest + +from hindsight_api.config import LLMMetadataRoute, LLMStrategyConfig, _get_raw_config +from hindsight_api.engine import memory_engine as engine_module +from hindsight_api.engine.llm_wrapper import LLMProvider +from hindsight_api.engine.multi_llm import MultiLLMProvider + + +@dataclasses.dataclass +class _CallCounts: + primary: int = 0 + secondary: int = 0 + + def reset(self) -> None: + self.primary = 0 + self.secondary = 0 + + +def _install_metadata_router( + memory, + monkeypatch, + *, + operation: str = "retain", + key: str = "tags", + value: str = "sensitive", +) -> _CallCounts: + counts = _CallCounts() + primary = LLMProvider(provider="mock", api_key="", base_url="", model="primary") + secondary = LLMProvider(provider="mock", api_key="", base_url="", model="secondary") + primary_call = primary.call + secondary_call = secondary.call + primary_call_with_tools = primary.call_with_tools + secondary_call_with_tools = secondary.call_with_tools + + async def record_primary(*args, **kwargs): + counts.primary += 1 + return await primary_call(*args, **kwargs) + + async def record_secondary(*args, **kwargs): + counts.secondary += 1 + return await secondary_call(*args, **kwargs) + + async def record_primary_with_tools(*args, **kwargs): + counts.primary += 1 + return await primary_call_with_tools(*args, **kwargs) + + async def record_secondary_with_tools(*args, **kwargs): + counts.secondary += 1 + return await secondary_call_with_tools(*args, **kwargs) + + monkeypatch.setattr(primary, "call", record_primary) + monkeypatch.setattr(secondary, "call", record_secondary) + monkeypatch.setattr(primary, "call_with_tools", record_primary_with_tools) + monkeypatch.setattr(secondary, "call_with_tools", record_secondary_with_tools) + setattr( + memory, + f"_{operation}_llm_config", + MultiLLMProvider( + [primary, secondary], + LLMStrategyConfig( + mode="metadata", + routes=[LLMMetadataRoute(key=key, value=value, member=1)], + ), + ), + ) + return counts + + +def _install_ambiguous_metadata_router(memory, monkeypatch) -> None: + members = [ + LLMProvider(provider="mock", api_key="", base_url="", model="primary"), + LLMProvider(provider="mock", api_key="", base_url="", model="internal"), + LLMProvider(provider="mock", api_key="", base_url="", model="sensitive"), + ] + + async def unexpected_call(*args, **kwargs): + pytest.fail("ambiguous retain reached an LLM provider") + + for member in members: + monkeypatch.setattr(member, "call", unexpected_call) + monkeypatch.setattr(member, "call_with_tools", unexpected_call) + + memory._retain_llm_config = MultiLLMProvider( + members, + LLMStrategyConfig( + mode="metadata", + routes=[ + LLMMetadataRoute(key="metadata.classification", value="internal", member=1), + LLMMetadataRoute(key="metadata.clearance", value="restricted", member=2), + ], + ), + ) + + +async def test_http_preserves_explicit_empty_classification(api_client, memory, monkeypatch) -> None: + captured: list[dict] = [] + + async def capture_retain(*args, **kwargs): + captured.extend(kwargs["contents"]) + return [[]], None + + monkeypatch.setattr(memory, "retain_batch_async", capture_retain) + response = await api_client.post( + "/v1/default/banks/metadata-routing-http/memories", + json={ + "items": [ + { + "content": "Explicitly declassified append.", + "document_id": "document-1", + "update_mode": "append", + "tags": [], + "metadata": {}, + } + ] + }, + ) + + assert response.status_code == 200 + assert captured[0]["tags"] == [] + assert captured[0]["metadata"] == {} + + +async def test_split_sync_retain_rejects_cross_member_classification_before_llm( + memory_no_llm_verify, request_context, monkeypatch +) -> None: + _install_ambiguous_metadata_router(memory_no_llm_verify, monkeypatch) + narrowed = dataclasses.replace(_get_raw_config(), retain_batch_tokens=20) + monkeypatch.setattr(engine_module, "get_config", lambda: narrowed) + bank_id = f"metadata-routing-mixed-sync-{uuid.uuid4().hex[:8]}" + contents = [ + {"content": " ".join(["internal"] * 80), "metadata": {"classification": "internal"}}, + {"content": " ".join(["sensitive"] * 80), "metadata": {"clearance": "restricted"}}, + ] + + with pytest.raises(ValueError, match="select multiple members"): + await memory_no_llm_verify.retain_batch_async( + bank_id, + contents, + request_context=request_context, + ) + + +async def test_queued_retain_rejects_cross_member_classification_before_children( + memory_no_llm_verify, request_context, monkeypatch +) -> None: + _install_ambiguous_metadata_router(memory_no_llm_verify, monkeypatch) + narrowed = dataclasses.replace(_get_raw_config(), retain_batch_tokens=20) + monkeypatch.setattr(engine_module, "get_config", lambda: narrowed) + bank_id = f"metadata-routing-mixed-queued-{uuid.uuid4().hex[:8]}" + contents = [ + {"content": " ".join(["internal"] * 80), "metadata": {"classification": "internal"}}, + {"content": " ".join(["sensitive"] * 80), "metadata": {"clearance": "restricted"}}, + ] + + with pytest.raises(ValueError, match="select multiple members"): + await memory_no_llm_verify.submit_async_retain( + bank_id, + contents, + request_context=request_context, + ) + + operations = await memory_no_llm_verify.list_operations(bank_id, request_context=request_context) + assert operations["total"] == 0 + + +async def test_tag_routed_reflect_stays_on_secondary_for_every_scope( + memory_no_llm_verify, request_context, monkeypatch +) -> None: + bank_id = f"metadata-routing-reflect-{uuid.uuid4().hex[:8]}" + await memory_no_llm_verify.retain_batch_async( + bank_id, + [{"content": "The restricted launch code is 2468.", "tags": ["sensitive"]}], + request_context=request_context, + ) + await memory_no_llm_verify.retain_batch_async( + bank_id, + [{"content": "The public office opens at nine.", "tags": ["public"]}], + request_context=request_context, + ) + + calls = _install_metadata_router(memory_no_llm_verify, monkeypatch, operation="reflect") + await memory_no_llm_verify.reflect_async( + bank_id, + "Summarize the available information.", + request_context=request_context, + ) + assert calls.secondary > 0 + assert calls.primary == 0 + + # Expand can turn a public fact into its full mixed-classification document, + # so even an exact public scope must stay on the protected lane. + calls.reset() + await memory_no_llm_verify.reflect_async( + bank_id, + "When does the office open?", + tags=["public"], + tags_match="exact", + request_context=request_context, + ) + assert calls.secondary > 0 + assert calls.primary == 0 + + calls.reset() + await memory_no_llm_verify.reflect_async( + bank_id, + "What is the launch code?", + tags=["sensitive"], + tags_match="exact", + request_context=request_context, + ) + assert calls.secondary > 0 + assert calls.primary == 0 + + +async def test_sensitive_pending_facts_route_consolidation_to_secondary( + memory_no_llm_verify, request_context, monkeypatch +) -> None: + from hindsight_api.engine.consolidation import consolidator + + calls = _install_metadata_router(memory_no_llm_verify, monkeypatch, operation="consolidation") + + async def capture_run(memory_engine, bank_id, context, config, llm_config, *args): + await llm_config.call(messages=[{"role": "user", "content": "Sensitive facts"}], scope="consolidation") + return {"status": "complete"} + + monkeypatch.setattr(consolidator, "_run_consolidation_job", capture_run) + result = await consolidator.run_consolidation_job( + memory_no_llm_verify, + "metadata-routing-consolidation", + request_context, + ) + + assert result == {"status": "complete"} + assert calls.secondary == 1 + assert calls.primary == 0 + + +async def test_sensitive_retain_uses_secondary_without_touching_primary( + memory_no_llm_verify, request_context, monkeypatch +) -> None: + calls = _install_metadata_router(memory_no_llm_verify, monkeypatch) + + bank_id = f"metadata-routing-{uuid.uuid4().hex[:8]}" + document_id = "private-account" + unit_ids = await memory_no_llm_verify.retain_batch_async( + bank_id, + [ + { + "content": "Alice's private account number is 1234.", + "document_id": document_id, + "tags": ["sensitive"], + } + ], + request_context=request_context, + ) + + assert unit_ids[0] + assert calls.secondary > 0 + assert calls.primary == 0 + + memories = await memory_no_llm_verify.list_memory_units( + bank_id, + fact_type=["world", "experience"], + tags=["sensitive"], + tags_match="all_strict", + request_context=request_context, + ) + assert memories["total"] == len(unit_ids[0]) + + calls.reset() + narrowed = dataclasses.replace(_get_raw_config(), retain_batch_tokens=20) + monkeypatch.setattr(engine_module, "get_config", lambda: narrowed) + sub_batch_count = 0 + real_iter_sub_batches = engine_module.iter_sub_batches + + def count_sub_batches(*args, **kwargs): + nonlocal sub_batch_count + for sub_batch in real_iter_sub_batches(*args, **kwargs): + sub_batch_count += 1 + yield sub_batch + + monkeypatch.setattr(engine_module, "iter_sub_batches", count_sub_batches) + await memory_no_llm_verify.retain_batch_async( + bank_id, + [ + { + "content": " ".join( + f"The private account review entry number {index} was completed today." for index in range(120) + ), + "document_id": document_id, + "update_mode": "append", + } + ], + request_context=request_context, + ) + + assert sub_batch_count > 1 + assert calls.secondary > 0 + assert calls.primary == 0 + + # Omitted tags inherit the stored classification, including for later + # appends after the first append has replaced the document record. + calls.reset() + await memory_no_llm_verify.retain_batch_async( + bank_id, + [{"content": "A second review was recorded.", "document_id": document_id, "update_mode": "append"}], + request_context=request_context, + ) + assert calls.secondary > 0 + assert calls.primary == 0 + document = await memory_no_llm_verify.get_document(document_id, bank_id, request_context=request_context) + assert document is not None and document["tags"] == ["sensitive"] + + # Adding an unrelated tag must not implicitly remove the classifier. + calls.reset() + await memory_no_llm_verify.retain_batch_async( + bank_id, + [ + { + "content": "Customer scope added.", + "document_id": document_id, + "update_mode": "append", + "tags": ["customer"], + } + ], + request_context=request_context, + ) + assert calls.secondary > 0 + assert calls.primary == 0 + document = await memory_no_llm_verify.get_document(document_id, bank_id, request_context=request_context) + assert document is not None and set(document["tags"]) == {"customer", "sensitive"} + + # Explicitly clearing tags declassifies future appends, but this operation + # still reprocesses the old sensitive body on the protected lane. + calls.reset() + await memory_no_llm_verify.retain_batch_async( + bank_id, + [{"content": "Classification cleared.", "document_id": document_id, "update_mode": "append", "tags": []}], + request_context=request_context, + ) + assert calls.secondary > 0 + assert calls.primary == 0 + + calls.reset() + await memory_no_llm_verify.retain_batch_async( + bank_id, + [{"content": "A public follow-up.", "document_id": document_id, "update_mode": "append"}], + request_context=request_context, + ) + assert calls.primary > 0 + assert calls.secondary == 0 + + +async def test_custom_metadata_is_inherited_for_append_routing( + memory_no_llm_verify, request_context, monkeypatch +) -> None: + calls = _install_metadata_router( + memory_no_llm_verify, + monkeypatch, + key="metadata.classification", + value="restricted", + ) + bank_id = f"metadata-routing-custom-{uuid.uuid4().hex[:8]}" + document_id = "restricted-document" + + await memory_no_llm_verify.retain_batch_async( + bank_id, + [ + { + "content": "Restricted source material.", + "document_id": document_id, + "metadata": {"classification": "restricted"}, + } + ], + request_context=request_context, + ) + assert calls.secondary > 0 + assert calls.primary == 0 + + for content in ("First append without metadata.", "Second append without metadata."): + calls.reset() + await memory_no_llm_verify.retain_batch_async( + bank_id, + [{"content": content, "document_id": document_id, "update_mode": "append"}], + request_context=request_context, + ) + assert calls.secondary > 0 + assert calls.primary == 0 + + # Adding an unrelated key retains the stored classifier. + calls.reset() + await memory_no_llm_verify.retain_batch_async( + bank_id, + [ + { + "content": "Source metadata added.", + "document_id": document_id, + "update_mode": "append", + "metadata": {"source": "crm"}, + } + ], + request_context=request_context, + ) + assert calls.secondary > 0 + assert calls.primary == 0 + document = await memory_no_llm_verify.get_document(document_id, bank_id, request_context=request_context) + assert document is not None + assert document["document_metadata"] == {"classification": "restricted", "source": "crm"} + + # An explicit empty map clears custom routing metadata for later appends. + calls.reset() + await memory_no_llm_verify.retain_batch_async( + bank_id, + [{"content": "Classification cleared.", "document_id": document_id, "update_mode": "append", "metadata": {}}], + request_context=request_context, + ) + assert calls.secondary > 0 + assert calls.primary == 0 + + calls.reset() + await memory_no_llm_verify.retain_batch_async( + bank_id, + [{"content": "Public follow-up.", "document_id": document_id, "update_mode": "append"}], + request_context=request_context, + ) + assert calls.primary > 0 + assert calls.secondary == 0 + + +async def test_shared_document_persists_non_first_item_metadata_for_append_routing( + memory_no_llm_verify, request_context, monkeypatch +) -> None: + calls = _install_metadata_router( + memory_no_llm_verify, + monkeypatch, + key="metadata.classification", + value="restricted", + ) + bank_id = f"metadata-routing-shared-{uuid.uuid4().hex[:8]}" + document_id = "shared-restricted-document" + contents = [ + {"content": "Public preface.", "document_id": document_id}, + { + "content": "Restricted details.", + "document_id": document_id, + "metadata": {"classification": "restricted"}, + }, + ] + original_contents = copy.deepcopy(contents) + + await memory_no_llm_verify.retain_batch_async( + bank_id, + contents, + request_context=request_context, + ) + assert calls.secondary > 0 + assert calls.primary == 0 + assert contents == original_contents + + document = await memory_no_llm_verify.get_document(document_id, bank_id, request_context=request_context) + assert document is not None + assert document["document_metadata"] == {"classification": "restricted"} + + calls.reset() + await memory_no_llm_verify.retain_batch_async( + bank_id, + [{"content": "Append without metadata.", "document_id": document_id, "update_mode": "append"}], + request_context=request_context, + ) + assert calls.secondary > 0 + assert calls.primary == 0 + + +async def test_append_routing_bulk_reads_store_after_releasing_sql_connection( + memory_no_llm_verify, monkeypatch +) -> None: + from hindsight_api.engine.memories import set_memories + from tests.test_memories_extension import InMemoryMemories + + connection_held = False + sql_calls: list[tuple[str, list[str], str]] = [] + + class FakeConnection: + async def fetch(self, query, document_ids, bank_id): + assert connection_held + sql_calls.append((query, document_ids, bank_id)) + return [] + + @asynccontextmanager + async def fake_acquire(_backend): + nonlocal connection_held + connection_held = True + try: + yield FakeConnection() + finally: + connection_held = False + + class BulkDocumentStore(InMemoryMemories): + def __init__(self): + super().__init__({}) + self.bulk_calls = 0 + + async def get_document_records(self, *, bank_id, document_ids): + assert not connection_held + self.bulk_calls += 1 + return { + "doc-a": {"tags": ["sensitive"], "metadata": {}}, + "doc-b": { + "tags": [], + "metadata": {"retain_params": {"metadata": {"classification": "restricted"}}}, + }, + } + + _install_metadata_router(memory_no_llm_verify, monkeypatch) + monkeypatch.setattr(engine_module, "acquire_with_retry", fake_acquire) + monkeypatch.setattr(memory_no_llm_verify, "_get_backend", AsyncMock(return_value=object())) + store = BulkDocumentStore() + set_memories(store) + try: + states = await memory_no_llm_verify._stored_append_routing_states( + "bank", + [ + {"content": "a", "document_id": "doc-a", "update_mode": "append"}, + {"content": "b", "document_id": "doc-b", "update_mode": "append"}, + ], + ) + finally: + set_memories(None) + + assert len(sql_calls) == 1 + assert "id = ANY($1::text[])" in sql_calls[0][0] + assert sql_calls[0][1:] == (["doc-a", "doc-b"], "bank") + assert store.bulk_calls == 1 + assert states["doc-a"].tags == ["sensitive"] + assert states["doc-b"].metadata == {"classification": "restricted"} + + +async def test_store_owned_sensitive_append_uses_authoritative_document_tags( + memory_no_llm_verify, request_context, monkeypatch +) -> None: + from hindsight_api.engine.memories import set_memories + from tests.test_memories_extension import InMemoryMemories + + class MetadataAwareStore(InMemoryMemories): + async def index_facts(self, bank_id, unit_ids, facts, document_id=None, unit_entity_ids=None): + # The store-owned retain session supplies its public FactRecord + # shape, while this shared test store predates that seam and still + # expects ProcessedFact. Preserve the store's observable behavior + # without making this routing regression depend on that mismatch. + from hindsight_api.engine.memories.base import StoredMemory + + self.calls.append("index_facts") + for unit_id, fact in zip(unit_ids, facts): + self.rows[unit_id] = StoredMemory( + unit_id=unit_id, + text=fact.text, + fact_type=fact.fact_type, + context=fact.context, + document_id=document_id, + chunk_id=fact.chunk_id, + tags=list(fact.tags or []), + metadata=fact.metadata, + created_at=fact.created_at, + ) + + async def get_document_record(self, *, bank_id, document_id, include_text=False): + record = await super().get_document_record( + bank_id=bank_id, + document_id=document_id, + include_text=include_text, + ) + if record is not None: + record["metadata"] = dict(self.documents[document_id]["metadata"]) + return record + + store = MetadataAwareStore({}) + set_memories(store) + try: + calls = _install_metadata_router(memory_no_llm_verify, monkeypatch) + bank_id = f"metadata-routing-store-{uuid.uuid4().hex[:8]}" + document_id = "store-owned-sensitive" + store.documents[document_id] = { + "id": document_id, + "content_hash": "sensitive-seed", + "original_text": "Sensitive store-owned content.", + "chunk_texts": ["Sensitive store-owned content."], + "chunks": ["Sensitive store-owned content."], + "tags": ["sensitive"], + "metadata": {"retain_params": {}}, + } + + calls.reset() + await memory_no_llm_verify.retain_batch_async( + bank_id, + [{"content": "An append without tags.", "document_id": document_id, "update_mode": "append"}], + request_context=request_context, + ) + assert calls.secondary > 0 + assert calls.primary == 0 + + calls = _install_metadata_router( + memory_no_llm_verify, + monkeypatch, + key="metadata.classification", + value="restricted", + ) + metadata_document_id = "store-owned-restricted" + store.documents[metadata_document_id] = { + "id": metadata_document_id, + "content_hash": "restricted-seed", + "original_text": "Restricted store-owned content.", + "chunk_texts": ["Restricted store-owned content."], + "chunks": ["Restricted store-owned content."], + "tags": [], + "metadata": {"retain_params": {"metadata": {"classification": "restricted"}}}, + } + + calls.reset() + await memory_no_llm_verify.retain_batch_async( + bank_id, + [ + { + "content": "An append without metadata.", + "document_id": metadata_document_id, + "update_mode": "append", + } + ], + request_context=request_context, + ) + assert calls.secondary > 0 + assert calls.primary == 0 + finally: + set_memories(None) diff --git a/hindsight-api-slim/tests/test_multi_llm_batch.py b/hindsight-api-slim/tests/test_multi_llm_batch.py index d6ec5dfe57..baf9eccc7b 100644 --- a/hindsight-api-slim/tests/test_multi_llm_batch.py +++ b/hindsight-api-slim/tests/test_multi_llm_batch.py @@ -15,7 +15,7 @@ import pytest -from hindsight_api.config import LLM_STRATEGY_FAILOVER, HindsightConfig, LLMStrategyConfig +from hindsight_api.config import LLM_STRATEGY_FAILOVER, HindsightConfig, LLMMetadataRoute, LLMStrategyConfig from hindsight_api.engine.multi_llm import MultiLLMProvider from hindsight_api.engine.retain.fact_extraction import RetainContent, extract_facts_from_contents_batch_api @@ -87,6 +87,9 @@ async def batch_provider_impl(self, account_key: str | None = None) -> _FakeBatc return None return self._provider_impl + def with_config(self, config: HindsightConfig, **kwargs: Any) -> "_BatchMember": + return self + class _FakeConn: """Serves the ``result_metadata`` read the resume path makes, and records the @@ -130,6 +133,16 @@ def _chain(*members: _BatchMember) -> MultiLLMProvider: return MultiLLMProvider(list(members), LLMStrategyConfig(mode=LLM_STRATEGY_FAILOVER)) +def _metadata_chain(*members: _BatchMember) -> MultiLLMProvider: + return MultiLLMProvider( + list(members), + LLMStrategyConfig( + mode="metadata", + routes=[LLMMetadataRoute(key="tags", value="sensitive", member=1)], + ), + ) + + def _batch_config() -> HindsightConfig: config = HindsightConfig.from_env() config.retain_batch_enabled = True @@ -169,6 +182,67 @@ async def test_batch_provider_impl_is_none_when_no_member_capable() -> None: assert await multi.batch_provider_impl() is None +async def test_metadata_routes_require_batch_support_on_every_selectable_member() -> None: + assert await _metadata_chain(_BatchMember("openai", True), _BatchMember("groq", True)).supports_batch_api() + assert not await _metadata_chain(_BatchMember("openai", True), _BatchMember("deepseek", False)).supports_batch_api() + + +async def test_metadata_batch_selection_defaults_to_primary_without_request_metadata() -> None: + primary = _BatchMember("deepseek", False) + secondary = _BatchMember("openai", True) + assert await _metadata_chain(primary, secondary).batch_provider_impl() is None + + +async def test_metadata_batch_lifecycle_runs_on_sensitive_member() -> None: + primary = _BatchMember("openai", True, account="primary") + secondary = _BatchMember("openai", True, account="sensitive") + configured = _metadata_chain(primary, secondary).with_config( + _batch_config(), + routing_metadata={"tags": ["sensitive"]}, + ) + + await extract_facts_from_contents_batch_api( + contents=[RetainContent(content="Sensitive account details.")], + llm_config=configured, + config=_batch_config(), + pool=None, + operation_id=None, + schema=None, + ) + + assert secondary._provider_impl.calls == ["submit", "status", "retrieve"] + assert primary._provider_impl.calls == [] + + +async def test_metadata_batch_resume_stays_on_sensitive_account() -> None: + primary = _BatchMember("openai", True, account="primary") + secondary = _BatchMember("openai", True, account="sensitive") + configured = _metadata_chain(primary, secondary).with_config( + _batch_config(), + routing_metadata={"tags": ["sensitive"]}, + ) + pool = _FakePool( + { + "batch_id": "batch_123", + "batch_provider": "openai", + "batch_account": secondary._provider_impl.batch_account_key, + "chunk_count": 1, + } + ) + + await extract_facts_from_contents_batch_api( + contents=[RetainContent(content="Sensitive account details.")], + llm_config=configured, + config=_batch_config(), + pool=pool, + operation_id=str(uuid.uuid4()), + schema=None, + ) + + assert secondary._provider_impl.calls == ["status", "retrieve"] + assert primary._provider_impl.calls == [] + + # ── the batch lifecycle targets the selected member ───────────────────────────── diff --git a/hindsight-api-slim/tests/test_multi_llm_config.py b/hindsight-api-slim/tests/test_multi_llm_config.py index ed03b90a64..1a37f17e55 100644 --- a/hindsight-api-slim/tests/test_multi_llm_config.py +++ b/hindsight-api-slim/tests/test_multi_llm_config.py @@ -10,12 +10,13 @@ from hindsight_api.config import ( HindsightConfig, LLMMemberConfig, + LLMMetadataRoute, LLMStrategyConfig, _parse_llm_members, _parse_llm_strategy, ) from hindsight_api.engine.llm_wrapper import LLMProvider -from hindsight_api.engine.memory_engine import _build_llm, _LLMCallDefaults +from hindsight_api.engine.memory_engine import _build_llm, _LLMCallDefaults, _retain_llm_routing_metadata from hindsight_api.engine.multi_llm import MultiLLMProvider # No per-request overrides — exercises the chain-resolution logic without @@ -206,6 +207,14 @@ def test_parse_strategy_weighted_round_robin(): assert s.weights == [3, 1] +def test_parse_strategy_metadata_routes(): + s = _parse_llm_strategy('{"mode": "metadata", "routes": [{"key": "tags", "value": "sensitive", "member": 1}]}') + assert s == LLMStrategyConfig( + mode="metadata", + routes=[LLMMetadataRoute(key="tags", value="sensitive", member=1)], + ) + + def test_parse_strategy_invalid_json(): with pytest.raises(ValueError, match="invalid JSON"): _parse_llm_strategy("{not json") @@ -228,6 +237,24 @@ def test_parse_strategy_weights_must_be_positive_ints(): _parse_llm_strategy('{"mode": "round-robin", "weights": []}') +@pytest.mark.parametrize( + "raw, message", + [ + ('{"mode": "metadata"}', "non-empty list"), + ('{"mode": "metadata", "routes": []}', "non-empty list"), + ('{"mode": "metadata", "routes": [{}]}', "'key'"), + ( + '{"mode": "metadata", "routes": [{"key": "tags", "value": "sensitive", "member": -1}]}', + "non-negative integer", + ), + ('{"mode": "failover", "routes": []}', "only valid with mode"), + ], +) +def test_parse_strategy_metadata_routes_validate(raw, message): + with pytest.raises(ValueError, match=message): + _parse_llm_strategy(raw) + + # ── from_env integration ──────────────────────────────────────────────────────── @@ -325,6 +352,32 @@ def test_build_llm_members_without_strategy_stays_plain(clean_llm_env): assert _build_llm(base, config, "", _NO_CALL_DEFAULTS) is base +def test_build_llm_rejects_metadata_route_to_missing_member(clean_llm_env): + config = _empty_config( + llm_members=[_member("ollama")], + llm_strategy=LLMStrategyConfig( + mode="metadata", + routes=[LLMMetadataRoute(key="tags", value="sensitive", member=2)], + ), + ) + with pytest.raises(ValueError, match="members 0..1"): + _build_llm(_base_llm(), config, "", _NO_CALL_DEFAULTS) + + +def test_retain_routing_metadata_unions_item_and_document_values(): + metadata = _retain_llm_routing_metadata( + [ + {"content": "public", "tags": ["public"], "metadata": {"classification": "internal"}}, + {"content": "secret", "tags": ["sensitive"], "metadata": {"classification": "restricted"}}, + ], + document_tags=["shared"], + ) + assert metadata == { + "metadata.classification": ["internal", "restricted"], + "tags": ["public", "sensitive", "shared"], + } + + # ── vertexai member build path (_member_to_llm) ───────────────────────────────── diff --git a/hindsight-api-slim/tests/test_multi_llm_provider.py b/hindsight-api-slim/tests/test_multi_llm_provider.py index 8acf7890ce..4408a4f27a 100644 --- a/hindsight-api-slim/tests/test_multi_llm_provider.py +++ b/hindsight-api-slim/tests/test_multi_llm_provider.py @@ -1,4 +1,4 @@ -"""Unit tests for MultiLLMProvider routing (failover + weighted round-robin). +"""Unit tests for MultiLLMProvider routing strategies. Deterministic: members are lightweight fakes that record calls and either return a sentinel or raise a chosen exception. No real providers / network. @@ -8,7 +8,7 @@ import pytest -from hindsight_api.config import LLMStrategyConfig +from hindsight_api.config import LLMMetadataRoute, LLMStrategyConfig from hindsight_api.engine.llm_wrapper import OutputTooLongError from hindsight_api.engine.multi_llm import ( MultiLLMProvider, @@ -51,6 +51,9 @@ async def verify_connection(self): async def cleanup(self): pass + def with_config(self, config, **kwargs): + return self + def _resolve(self): b = self.behavior if isinstance(b, BaseException): @@ -68,6 +71,16 @@ def _round_robin(*members, weights=None): return MultiLLMProvider(list(members), LLMStrategyConfig(mode="round-robin", weights=weights)) +def _metadata(*members): + return MultiLLMProvider( + list(members), + LLMStrategyConfig( + mode="metadata", + routes=[LLMMetadataRoute(key="tags", value="sensitive", member=1)], + ), + ) + + # ── failover ───────────────────────────────────────────────────────────────── @@ -161,6 +174,88 @@ async def test_weighted_round_robin_honors_ratio(): assert (a.calls, b.calls) == (6, 2) +# ── metadata routing ────────────────────────────────────────────────────────── + + +async def test_metadata_route_pins_matching_tag_to_secondary(): + a, b = FakeMember("a", "RA"), FakeMember("b", "RB") + routed = _metadata(a, b).with_config(object(), routing_metadata={"tags": ["sensitive"]}) + assert await routed.call(messages=[]) == "RB" + assert (a.calls, b.calls) == (0, 1) + + +async def test_metadata_route_defaults_to_primary(): + a, b = FakeMember("a", "RA"), FakeMember("b", "RB") + routed = _metadata(a, b).with_config(object(), routing_metadata={"tags": ["public"]}) + assert await routed.call(messages=[]) == "RA" + assert (a.calls, b.calls) == (1, 0) + + +async def test_metadata_route_does_not_fall_back_across_privacy_lanes(): + a = FakeMember("a", "RA") + b = FakeMember("b", RuntimeError("sensitive lane down")) + routed = _metadata(a, b).with_config(object(), routing_metadata={"tags": ["sensitive"]}) + with pytest.raises(RuntimeError, match="sensitive lane down"): + await routed.call(messages=[]) + assert (a.calls, b.calls) == (0, 1) + + +def test_metadata_route_rejects_ambiguous_cross_member_match(): + a, b, c = FakeMember("a", "RA"), FakeMember("b", "RB"), FakeMember("c", "RC") + router = MultiLLMProvider( + [a, b, c], + LLMStrategyConfig( + mode="metadata", + routes=[ + LLMMetadataRoute(key="metadata.classification", value="internal", member=1), + LLMMetadataRoute(key="metadata.clearance", value="restricted", member=2), + ], + ), + ) + + with pytest.raises(ValueError, match="select multiple members"): + router.with_config( + object(), + routing_metadata={ + "metadata.classification": ["internal"], + "metadata.clearance": ["restricted"], + }, + ) + assert (a.calls, b.calls, c.calls) == (0, 0, 0) + + +def test_metadata_routes_reject_conflicting_tag_members_at_construction(): + with pytest.raises(ValueError, match="routes with key 'tags' must all select the same member"): + MultiLLMProvider( + [FakeMember("a", "RA"), FakeMember("b", "RB"), FakeMember("c", "RC")], + LLMStrategyConfig( + mode="metadata", + routes=[ + LLMMetadataRoute(key="tags", value="internal", member=1), + LLMMetadataRoute(key="tags", value="sensitive", member=2), + ], + ), + ) + + +async def test_metadata_route_allows_multiple_matches_to_same_member(): + a, b = FakeMember("a", "RA"), FakeMember("b", "RB") + router = MultiLLMProvider( + [a, b], + LLMStrategyConfig( + mode="metadata", + routes=[ + LLMMetadataRoute(key="tags", value="internal", member=1), + LLMMetadataRoute(key="tags", value="sensitive", member=1), + ], + ), + ) + + routed = router.with_config(object(), routing_metadata={"tags": ["internal", "sensitive"]}) + assert await routed.call(messages=[]) == "RB" + assert (a.calls, b.calls) == (0, 1) + + def test_weighted_scheduler_distribution(): sched = _WeightedRoundRobin([5, 1]) picks = [sched.next() for _ in range(6)] diff --git a/hindsight-api-slim/tests/test_retain_params_roundtrip.py b/hindsight-api-slim/tests/test_retain_params_roundtrip.py index 190c012674..3f9a75b16f 100644 --- a/hindsight-api-slim/tests/test_retain_params_roundtrip.py +++ b/hindsight-api-slim/tests/test_retain_params_roundtrip.py @@ -65,6 +65,20 @@ def test_a_new_field_round_trips_without_being_listed(): assert _params(content="x", some_future_field="v")["some_future_field"] == "v" +def test_shared_document_metadata_unions_values_from_every_item(): + retain_params, _tags = orch._build_retain_params( + [ + {"content": "public", "metadata": {"classification": "public", "source": "crm"}}, + {"content": "restricted", "metadata": {"classification": "restricted"}}, + ] + ) + + assert retain_params["metadata"] == { + "classification": ["public", "restricted"], + "source": "crm", + } + + def test_every_field_api_retain_sends_can_round_trip(): """Pairs the writer against the rule so they cannot drift apart. diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index 1349a13405..4af9e19cea 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -526,11 +526,11 @@ export HINDSIGHT_API_LLM_LITELLMROUTER_CONFIG='{ The config is a credential field — never returned by the bank-config API. Hindsight already retries calls; set `"num_retries": 0` in the Router config to avoid double-retries. Batch APIs aren't supported in router mode. -### Multi-LLM Strategies (failover / round-robin) +### Multi-LLM Strategies (failover / round-robin / metadata) Configure additional LLMs **by index** alongside the primary, then choose a strategy for routing across them. This is a provider-agnostic alternative to the LiteLLM Router: the indexed LLMs can be any mix of providers, each fully configured. -The unindexed `HINDSIGHT_API_LLM_*` config is the **primary** (member 1). Extra members are numbered from 1: +The unindexed `HINDSIGHT_API_LLM_*` config is the **primary** (routing member `0`). Extra members are numbered from `1`, and their environment suffix is also their routing member index: | Variable | Description | Default | |----------|-------------|---------| @@ -546,10 +546,11 @@ The unindexed `HINDSIGHT_API_LLM_*` config is the **primary** (member 1). Extra | `HINDSIGHT_API_LLM__LITELLMROUTER_CONFIG` | Per-member LiteLLM Router config JSON (for a `litellmrouter` member). Falls back to the global `HINDSIGHT_API_LLM_LITELLMROUTER_CONFIG` when unset. | - | | `HINDSIGHT_API_LLM_STRATEGY` | JSON routing strategy across the chain. Unset = single primary LLM (no change). | - | -The strategy JSON supports two modes: +The strategy JSON supports three modes: - `{"mode": "failover"}` — try members in order (primary first); on a member's failure (after its own retries) advance to the next. - `{"mode": "round-robin"}` — rotate the starting member per request to spread load, then fall through the rest on failure. Add `"weights": [3, 1, ...]` (positive ints, one per member, primary first) for an **unbalanced** rotation. +- `{"mode": "metadata", "routes": [...]}` — evaluate operation-metadata routes and pin the operation to its matching member. An unmatched operation uses member `0` (the primary). A matched operation does **not** fall back to another member if its selected model fails, so sensitive data cannot cross into a differently classified provider lane. ```bash # Primary OpenAI, failover to Groq then Anthropic @@ -563,11 +564,31 @@ export HINDSIGHT_API_LLM_STRATEGY='{"mode": "failover"}' # Weighted round-robin: serve the primary 3x as often as member 1 export HINDSIGHT_API_LLM_STRATEGY='{"mode": "round-robin", "weights": [3, 1]}' + +# Keep sensitive memories on member 1 across retain, reflect, and consolidation +export HINDSIGHT_API_LLM_1_PROVIDER=ollama +export HINDSIGHT_API_LLM_1_MODEL=qwen3:8b +export HINDSIGHT_API_LLM_STRATEGY='{ + "mode": "metadata", + "routes": [{"key": "tags", "value": "sensitive", "member": 1}] +}' ``` +Each metadata route has a string `key`, a string `value`, and a non-negative `member` index. Use `tags` to match an item or document tag, as above. User-defined retain metadata is exposed to the router under a `metadata.` prefix, so `{"key": "metadata.classification", "value": "restricted", "member": 1}` matches an item submitted with `"metadata": {"classification": "restricted"}`. Routing values are ephemeral and are not added to LLM trace metadata. + +Retain can combine several items in one LLM prompt. Hindsight unions their tags and metadata values before routing and persists the union for shared documents, so if any item matches a route, the entire combined prompt and later appends use that member. Multiple matching routes may target the same member. If one operation matches routes to different members, Hindsight rejects it rather than disclosing either classification to the other member; split the input or route those classifications to one member. Because reflect and consolidation must include every configured tag route to stay fail-closed, all `tags` routes must select the same member; conflicting tag routes are rejected at startup. + +Append reprocesses the existing document body, so Hindsight also routes using the document's stored tags and metadata. Stored classification is inherited when an append omits those fields, and merged into nonempty updates so adding an unrelated tag or metadata key does not accidentally remove it. A supplied metadata value replaces the stored value for the same key. Supplying an explicit empty collection (`"tags": []` or `"metadata": {}`) clears that classification for future operations; the clearing append itself still uses the previous route because its prompt includes the previously classified body. + +Retain routes from item tags and user-defined metadata. Reflect selects its model before the agent retrieves data, and its expand tool can return a fact's full document. Every reflect therefore includes all configured tag routes and uses their shared member, even for an exact public-only scope. This fail-closed behavior keeps a mixed-classification document or concurrent sensitive retain from entering a review after it has already selected the primary. + +Consolidation also selects one model for a job that can repeatedly fetch newly retained facts. With tag routes configured, it therefore includes all configured route tags and pins the job to their shared member. Mental-model refresh uses the same fail-closed reflect routing, including its delta-operation calls. + +Only retain has user-defined `metadata.*` values. Reflect, mental-model refresh, and consolidation route from stored tags. Configure the strategy globally, as above, to apply it to each operation, or repeat the members and strategy under the `RETAIN`, `REFLECT`, and `CONSOLIDATION` prefixes when those operations need different providers. + **Per-operation chains.** Each operation can define its own members + strategy with the `RETAIN` / `REFLECT` / `CONSOLIDATION` prefix (e.g. `HINDSIGHT_API_RETAIN_LLM_1_PROVIDER`, `HINDSIGHT_API_RETAIN_LLM_STRATEGY`). A per-operation slot with no indexed members (or no strategy) inherits the global chain. -The indexed members are credential fields — never returned by the bank-config API and server-level only (not per-bank configurable). **Batch retain** runs on the first batch-capable member in declared order, which need not be the primary — so a chain whose primary has no batch API can still use `HINDSIGHT_API_RETAIN_BATCH_ENABLED=true` as long as one member supports it. That member serves the whole batch (submit, polling and retrieval all target the account that holds it), so batch does not fail over the way the interactive retain/reflect/consolidation calls do. An in-flight batch is bound to the account that submitted it, so if the worker restarts mid-batch it resumes on that same account even when the chain has since been reordered or extended. Removing that member — or rotating its API key — while a batch is still running makes the operation fail with an explicit error instead of polling a different account. +The indexed members are credential fields — never returned by the bank-config API and server-level only (not per-bank configurable). **Batch retain** with failover or round-robin runs on the first batch-capable member in declared order, which need not be the primary — so a chain whose primary has no batch API can still use `HINDSIGHT_API_RETAIN_BATCH_ENABLED=true` as long as one member supports it. Metadata routing instead requires the primary and every route-selectable member to support batch, because each operation is pinned before submission. The selected member serves the whole batch (submit, polling and retrieval all target the account that holds it), so batch does not fail over the way the interactive failover/round-robin calls do. An in-flight batch is bound to the account that submitted it, so if the worker restarts mid-batch it resumes on that same account even when the chain has since been reordered or extended. Removing that member — or rotating its API key — while a batch is still running makes the operation fail with an explicit error instead of polling a different account. ### Built-in llama.cpp diff --git a/hindsight-embed/hindsight_embed/env.example b/hindsight-embed/hindsight_embed/env.example index df447edcb7..c8dc867de4 100644 --- a/hindsight-embed/hindsight_embed/env.example +++ b/hindsight-embed/hindsight_embed/env.example @@ -165,8 +165,11 @@ HINDSIGHT_API_LLM_MODEL=gpt-4o-mini # two Codex members fails over between two ChatGPT accounts instead of retrying one. # HINDSIGHT_API_LLM_1_PROVIDER=openai-codex # HINDSIGHT_API_LLM_1_CODEX_HOME=/var/lib/hindsight/codex-b -# Strategy JSON: {"mode": "failover"} or {"mode": "round-robin"}. +# Strategy JSON: {"mode": "failover"}, {"mode": "round-robin"}, or +# {"mode": "metadata", "routes": [{"key": "tags", "value": "sensitive", "member": 1}]}. # Round-robin accepts optional positive-int "weights" (one per member, primary first). +# Metadata routes are strict: all matches must select one member. An unmatched +# request uses member 0 (the primary); a match never falls across lanes. # HINDSIGHT_API_LLM_STRATEGY={"mode": "failover"} # API Configuration (Optional) diff --git a/skills/hindsight-docs/references/developer/configuration.md b/skills/hindsight-docs/references/developer/configuration.md index 6de14a976d..b0c161c211 100644 --- a/skills/hindsight-docs/references/developer/configuration.md +++ b/skills/hindsight-docs/references/developer/configuration.md @@ -526,11 +526,11 @@ export HINDSIGHT_API_LLM_LITELLMROUTER_CONFIG='{ The config is a credential field — never returned by the bank-config API. Hindsight already retries calls; set `"num_retries": 0` in the Router config to avoid double-retries. Batch APIs aren't supported in router mode. -### Multi-LLM Strategies (failover / round-robin) +### Multi-LLM Strategies (failover / round-robin / metadata) Configure additional LLMs **by index** alongside the primary, then choose a strategy for routing across them. This is a provider-agnostic alternative to the LiteLLM Router: the indexed LLMs can be any mix of providers, each fully configured. -The unindexed `HINDSIGHT_API_LLM_*` config is the **primary** (member 1). Extra members are numbered from 1: +The unindexed `HINDSIGHT_API_LLM_*` config is the **primary** (routing member `0`). Extra members are numbered from `1`, and their environment suffix is also their routing member index: | Variable | Description | Default | |----------|-------------|---------| @@ -546,10 +546,11 @@ The unindexed `HINDSIGHT_API_LLM_*` config is the **primary** (member 1). Extra | `HINDSIGHT_API_LLM__LITELLMROUTER_CONFIG` | Per-member LiteLLM Router config JSON (for a `litellmrouter` member). Falls back to the global `HINDSIGHT_API_LLM_LITELLMROUTER_CONFIG` when unset. | - | | `HINDSIGHT_API_LLM_STRATEGY` | JSON routing strategy across the chain. Unset = single primary LLM (no change). | - | -The strategy JSON supports two modes: +The strategy JSON supports three modes: - `{"mode": "failover"}` — try members in order (primary first); on a member's failure (after its own retries) advance to the next. - `{"mode": "round-robin"}` — rotate the starting member per request to spread load, then fall through the rest on failure. Add `"weights": [3, 1, ...]` (positive ints, one per member, primary first) for an **unbalanced** rotation. +- `{"mode": "metadata", "routes": [...]}` — evaluate operation-metadata routes and pin the operation to its matching member. An unmatched operation uses member `0` (the primary). A matched operation does **not** fall back to another member if its selected model fails, so sensitive data cannot cross into a differently classified provider lane. ```bash # Primary OpenAI, failover to Groq then Anthropic @@ -563,11 +564,31 @@ export HINDSIGHT_API_LLM_STRATEGY='{"mode": "failover"}' # Weighted round-robin: serve the primary 3x as often as member 1 export HINDSIGHT_API_LLM_STRATEGY='{"mode": "round-robin", "weights": [3, 1]}' + +# Keep sensitive memories on member 1 across retain, reflect, and consolidation +export HINDSIGHT_API_LLM_1_PROVIDER=ollama +export HINDSIGHT_API_LLM_1_MODEL=qwen3:8b +export HINDSIGHT_API_LLM_STRATEGY='{ + "mode": "metadata", + "routes": [{"key": "tags", "value": "sensitive", "member": 1}] +}' ``` +Each metadata route has a string `key`, a string `value`, and a non-negative `member` index. Use `tags` to match an item or document tag, as above. User-defined retain metadata is exposed to the router under a `metadata.` prefix, so `{"key": "metadata.classification", "value": "restricted", "member": 1}` matches an item submitted with `"metadata": {"classification": "restricted"}`. Routing values are ephemeral and are not added to LLM trace metadata. + +Retain can combine several items in one LLM prompt. Hindsight unions their tags and metadata values before routing and persists the union for shared documents, so if any item matches a route, the entire combined prompt and later appends use that member. Multiple matching routes may target the same member. If one operation matches routes to different members, Hindsight rejects it rather than disclosing either classification to the other member; split the input or route those classifications to one member. Because reflect and consolidation must include every configured tag route to stay fail-closed, all `tags` routes must select the same member; conflicting tag routes are rejected at startup. + +Append reprocesses the existing document body, so Hindsight also routes using the document's stored tags and metadata. Stored classification is inherited when an append omits those fields, and merged into nonempty updates so adding an unrelated tag or metadata key does not accidentally remove it. A supplied metadata value replaces the stored value for the same key. Supplying an explicit empty collection (`"tags": []` or `"metadata": {}`) clears that classification for future operations; the clearing append itself still uses the previous route because its prompt includes the previously classified body. + +Retain routes from item tags and user-defined metadata. Reflect selects its model before the agent retrieves data, and its expand tool can return a fact's full document. Every reflect therefore includes all configured tag routes and uses their shared member, even for an exact public-only scope. This fail-closed behavior keeps a mixed-classification document or concurrent sensitive retain from entering a review after it has already selected the primary. + +Consolidation also selects one model for a job that can repeatedly fetch newly retained facts. With tag routes configured, it therefore includes all configured route tags and pins the job to their shared member. Mental-model refresh uses the same fail-closed reflect routing, including its delta-operation calls. + +Only retain has user-defined `metadata.*` values. Reflect, mental-model refresh, and consolidation route from stored tags. Configure the strategy globally, as above, to apply it to each operation, or repeat the members and strategy under the `RETAIN`, `REFLECT`, and `CONSOLIDATION` prefixes when those operations need different providers. + **Per-operation chains.** Each operation can define its own members + strategy with the `RETAIN` / `REFLECT` / `CONSOLIDATION` prefix (e.g. `HINDSIGHT_API_RETAIN_LLM_1_PROVIDER`, `HINDSIGHT_API_RETAIN_LLM_STRATEGY`). A per-operation slot with no indexed members (or no strategy) inherits the global chain. -The indexed members are credential fields — never returned by the bank-config API and server-level only (not per-bank configurable). **Batch retain** runs on the first batch-capable member in declared order, which need not be the primary — so a chain whose primary has no batch API can still use `HINDSIGHT_API_RETAIN_BATCH_ENABLED=true` as long as one member supports it. That member serves the whole batch (submit, polling and retrieval all target the account that holds it), so batch does not fail over the way the interactive retain/reflect/consolidation calls do. An in-flight batch is bound to the account that submitted it, so if the worker restarts mid-batch it resumes on that same account even when the chain has since been reordered or extended. Removing that member — or rotating its API key — while a batch is still running makes the operation fail with an explicit error instead of polling a different account. +The indexed members are credential fields — never returned by the bank-config API and server-level only (not per-bank configurable). **Batch retain** with failover or round-robin runs on the first batch-capable member in declared order, which need not be the primary — so a chain whose primary has no batch API can still use `HINDSIGHT_API_RETAIN_BATCH_ENABLED=true` as long as one member supports it. Metadata routing instead requires the primary and every route-selectable member to support batch, because each operation is pinned before submission. The selected member serves the whole batch (submit, polling and retrieval all target the account that holds it), so batch does not fail over the way the interactive failover/round-robin calls do. An in-flight batch is bound to the account that submitted it, so if the worker restarts mid-batch it resumes on that same account even when the chain has since been reordered or extended. Removing that member — or rotating its API key — while a batch is still running makes the operation fail with an explicit error instead of polling a different account. ### Built-in llama.cpp