claude-smart 0.2.49 → 0.2.50

package/plugin/vendor/reflexio/reflexio/server/services/playbook/components/consolidator.py

@@ -13,7 +13,6 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator
 from reflexio.models.api_schema.retriever_schema import SearchUserPlaybookRequest
 from reflexio.models.api_schema.service_schemas import UserPlaybook
 from reflexio.models.config_schema import (
-    EMBEDDING_DIMENSIONS,
     DeduplicationConfig,
     SearchOptions,
 )
@@ -23,6 +22,10 @@ from reflexio.server.llm.litellm_client import LiteLLMClient
 from reflexio.server.services.deduplication_utils import (
     BaseDeduplicator,
     format_dedup_timestamp,
+    resolve_dedup_query_embeddings,
+)
+from reflexio.server.services.profile.profile_generation_service_utils import (
+    check_string_token_overlap,
 )
 from reflexio.server.tracing import sentry_tags
 
@@ -191,6 +194,22 @@ def _existing_id_log_label(value: ExistingIdField) -> str:
     return str(value)
 
 
+def _normalize_generation_request_id(
+    generation_request_id: str | None,
+    *,
+    request_id: str | None = None,
+) -> str:
+    if generation_request_id is not None:
+        if request_id is not None and request_id != generation_request_id:
+            raise TypeError(
+                "generation_request_id and request_id must match when both are provided"
+            )
+        return generation_request_id
+    if request_id is not None:
+        return request_id
+    raise TypeError("generation_request_id is required")
+
+
 class UnifyDecision(BaseModel):
     """Collapse NEW (+ 0..N EXISTING) into one row with LLM-supplied content.
 
@@ -379,6 +398,50 @@ _COUNTER_BY_KIND: dict[str, str] = {
 }
 
 
+def _decision_new_ids(decision: ConsolidationDecision) -> list[str]:
+    if isinstance(decision, DifferentiateDecision):
+        return [decision.new_id]
+    return decision.new_ids
+
+
+def validate_consolidation_output(
+    new_playbooks: list[UserPlaybook],
+    output: PlaybookConsolidationOutput,
+) -> list[str]:
+    """Return partition-contract errors for a parsed consolidation output."""
+    known_new_ids = {f"NEW-{idx}" for idx in range(len(new_playbooks))}
+    seen_by_id: dict[str, list[str]] = {}
+    errors: list[str] = []
+
+    for decision_index, decision in enumerate(output.decisions):
+        label = f"decision[{decision_index}] {decision.kind}"
+        for new_id in _decision_new_ids(decision):
+            if new_id not in known_new_ids:
+                errors.append(f"{label} references unknown NEW id {new_id}.")
+                continue
+            seen_by_id.setdefault(new_id, []).append(label)
+
+    missing = sorted(known_new_ids - set(seen_by_id))
+    if missing:
+        errors.append(
+            "Every NEW id must appear exactly once; missing NEW ids: "
+            + ", ".join(missing)
+            + "."
+        )
+
+    duplicates = {
+        new_id: labels for new_id, labels in seen_by_id.items() if len(labels) > 1
+    }
+    for new_id, labels in sorted(duplicates.items()):
+        errors.append(
+            f"Every NEW id must appear exactly once; {new_id} appears in "
+            + "; ".join(labels)
+            + "."
+        )
+
+    return errors
+
+
 class PlaybookConsolidator(BaseDeduplicator):
     """
     Consolidates new user playbook entries against each other and against existing entries
@@ -498,15 +561,12 @@ class PlaybookConsolidator(BaseDeduplicator):
         if not query_texts:
             return []
 
-        # Batch-generate embeddings
-        try:
-            embeddings = self.client.get_embeddings(
-                query_texts, dimensions=EMBEDDING_DIMENSIONS
-            )
-        except Exception as e:
-            logger.warning("Failed to generate embeddings for dedup search: %s", e)
-            # Fall back to text-only search
-            embeddings = [None] * len(query_texts)
+        # Embed dedup queries with the same model that indexed the store —
+        # see resolve_dedup_query_embeddings for why the client's default
+        # embedding model must not be used here.
+        embeddings = resolve_dedup_query_embeddings(
+            storage, self.client, query_texts, entity_label="Playbook"
+        )
 
         # Search for each new entry
         seen_ids: set[int] = set()
@@ -565,6 +625,36 @@ class PlaybookConsolidator(BaseDeduplicator):
         )
         return new_text, existing_text
 
+    def _render_consolidation_prompt(
+        self,
+        new_playbooks: list[UserPlaybook],
+        existing_playbooks: list[UserPlaybook],
+    ) -> str:
+        """Render the consolidation prompt for the given NEW + EXISTING rows.
+
+        Rendering is deterministic, so the repair path calls this again to
+        reconstruct the exact first-turn prompt when building the follow-up
+        conversation.
+
+        Args:
+            new_playbooks: New entries to deduplicate.
+            existing_playbooks: Existing entries from the database.
+
+        Returns:
+            The fully rendered consolidation prompt.
+        """
+        new_text, existing_text = self._format_new_and_existing_for_prompt(
+            new_playbooks, existing_playbooks
+        )
+        return self.request_context.prompt_manager.render_prompt(
+            self._get_prompt_id(),
+            {
+                "new_playbook_count": len(new_playbooks),
+                "new_playbooks": new_text,
+                "existing_playbooks": existing_text,
+            },
+        )
+
     def _consolidation_decisions(
         self,
         new_playbooks: list[UserPlaybook],
@@ -605,20 +695,7 @@ class PlaybookConsolidator(BaseDeduplicator):
             Parsed ``PlaybookConsolidationOutput``; an empty output (no
             decisions) if the LLM returned an unexpected response shape.
         """
-        # Format for prompt
-        new_text, existing_text = self._format_new_and_existing_for_prompt(
-            new_playbooks, existing_playbooks
-        )
-
-        # Build and call LLM
-        prompt = self.request_context.prompt_manager.render_prompt(
-            self._get_prompt_id(),
-            {
-                "new_playbook_count": len(new_playbooks),
-                "new_playbooks": new_text,
-                "existing_playbooks": existing_text,
-            },
-        )
+        prompt = self._render_consolidation_prompt(new_playbooks, existing_playbooks)
 
         output_schema_class = self._get_output_schema_class()
 
@@ -653,16 +730,18 @@ class PlaybookConsolidator(BaseDeduplicator):
     def deduplicate(
         self,
         results: list[list[UserPlaybook]],
-        request_id: str,
-        agent_version: str,
+        generation_request_id: str | None = None,
+        agent_version: str | None = None,
         user_id: str | None = None,
+        *,
+        request_id: str | None = None,
     ) -> tuple[list[UserPlaybook], list[int], list[tuple[int, list[int]]]]:
         """
         Consolidate user playbook entries across extractors and against existing entries in DB.
 
         Args:
             results: List of entry lists from extractors (each extractor returns list[UserPlaybook])
-            request_id: Request ID for context
+            generation_request_id: Request ID for context
             agent_version: Agent version for context
             user_id: Optional user ID to scope the existing entry search
 
@@ -680,6 +759,12 @@ class PlaybookConsolidator(BaseDeduplicator):
             non-merge archives such as ``differentiate``'s split source); the
             caller subtracts the merge-covered ids to find pure-delete leftovers.
         """
+        generation_request_id = _normalize_generation_request_id(
+            generation_request_id, request_id=request_id
+        )
+        if agent_version is None:
+            raise TypeError("agent_version is required")
+
         # Check if mock mode is enabled
         if os.getenv("MOCK_LLM_RESPONSE", "").lower() == "true":
             logger.info("Mock mode: skipping consolidation")
@@ -714,7 +799,7 @@ class PlaybookConsolidator(BaseDeduplicator):
                 op="identify_duplicates",
                 org_id=self.request_context.org_id,
                 user_id=user_id,
-                request_id=request_id,
+                request_id=generation_request_id,
                 agent_version=agent_version,
                 error_type=type(e).__name__,
             ):
@@ -723,14 +808,21 @@ class PlaybookConsolidator(BaseDeduplicator):
 
         if not dedup_output.decisions:
             logger.info(
-                "No consolidation decisions returned for request %s", request_id
+                "No consolidation decisions returned for request %s",
+                generation_request_id,
             )
             return new_playbooks, [], []
 
         logger.info(
             "Received %d consolidation decisions for request %s",
             len(dedup_output.decisions),
-            request_id,
+            generation_request_id,
+        )
+        dedup_output = self._repair_consolidation_output_if_needed(
+            new_playbooks=new_playbooks,
+            existing_playbooks=existing_playbooks,
+            dedup_output=dedup_output,
+            request_id=request_id,
         )
 
         # Build consolidated result via the discriminated-union apply path
@@ -738,10 +830,215 @@ class PlaybookConsolidator(BaseDeduplicator):
             new_playbooks=new_playbooks,
             existing_playbooks=existing_playbooks,
             dedup_output=dedup_output,
-            request_id=request_id,
+            generation_request_id=generation_request_id,
             agent_version=agent_version,
         )
 
+    def _repair_consolidation_output_if_needed(
+        self,
+        *,
+        new_playbooks: list[UserPlaybook],
+        existing_playbooks: list[UserPlaybook],
+        dedup_output: PlaybookConsolidationOutput,
+        request_id: str | None,
+    ) -> PlaybookConsolidationOutput:
+        errors = validate_consolidation_output(new_playbooks, dedup_output)
+        if not errors:
+            errors = self._find_suspicious_under_consumed_new_rows(
+                new_playbooks, dedup_output
+            )
+        if not errors:
+            return dedup_output
+
+        repaired = self._repair_consolidation_output(
+            original_prompt=self._render_consolidation_prompt(
+                new_playbooks, existing_playbooks
+            ),
+            invalid_output=dedup_output,
+            errors=errors,
+            request_id=request_id,
+        )
+        if repaired is None:
+            return dedup_output
+
+        repaired_errors = validate_consolidation_output(new_playbooks, repaired)
+        if repaired_errors:
+            logger.warning(
+                "event=consolidation_repair_failed request_id=%s errors=%s",
+                request_id,
+                repaired_errors,
+            )
+            return dedup_output
+
+        logger.info(
+            "event=consolidation_repaired request_id=%s decision_count=%d",
+            request_id,
+            len(repaired.decisions),
+        )
+        return repaired
+
+    def _repair_consolidation_output(
+        self,
+        *,
+        original_prompt: str,
+        invalid_output: PlaybookConsolidationOutput,
+        errors: list[str],
+        request_id: str | None,
+    ) -> PlaybookConsolidationOutput | None:
+        """Ask the model to fix its consolidation output in a follow-up turn.
+
+        The repair is sent as the third turn of the original conversation —
+        [user: consolidation prompt, assistant: the invalid decisions, user:
+        validation errors + fix instruction] — so the model retains the full
+        first-turn context (rendered rows, decision-kind rules, output format)
+        without restating any of it.
+
+        Args:
+            original_prompt: The exact rendered first-turn consolidation prompt.
+            invalid_output: The parsed decisions that failed validation.
+            errors: Human-readable validation errors to feed back.
+            request_id: Request ID for log correlation.
+
+        Returns:
+            The repaired ``PlaybookConsolidationOutput``, or None when the
+            repair call fails or returns an unexpected shape.
+        """
+        logger.warning(
+            "event=consolidation_repair_attempted request_id=%s errors=%s",
+            request_id,
+            errors,
+        )
+        followup = "\n".join(
+            [
+                "Your decisions above violate the NEW-id coverage contract:",
+                "",
+                "\n".join(f"- {error}" for error in errors),
+                "",
+                "Return corrected decisions that cover every NEW id exactly once. "
+                "If a unified rule uses facts/details from multiple NEW rows, `new_id` must list all of them. "
+                "If your decisions are semantically correct apart from these errors, "
+                "keep them and fix only the coverage.",
+                "",
+                'Respond ONLY with valid JSON matching PlaybookConsolidationOutput: {"decisions": [...]}.',
+            ]
+        )
+        messages = [
+            {"role": "user", "content": original_prompt},
+            {"role": "assistant", "content": invalid_output.model_dump_json()},
+            {"role": "user", "content": followup},
+        ]
+
+        from reflexio.server.services.service_utils import (
+            log_llm_messages,
+            log_model_response,
+        )
+
+        log_llm_messages(logger, "Playbook consolidation repair", messages)
+
+        try:
+            response = self.client.generate_chat_response(
+                messages=messages,
+                model=self.model_name,
+                response_format=self._get_output_schema_class(),
+            )
+        except Exception as exc:  # noqa: BLE001 — repair is best-effort
+            logger.warning(
+                "event=consolidation_repair_failed request_id=%s error_type=%s",
+                request_id,
+                type(exc).__name__,
+                exc_info=True,
+            )
+            return None
+
+        log_model_response(logger, "Consolidation repair response", response)
+
+        if not isinstance(response, PlaybookConsolidationOutput):
+            logger.warning(
+                "event=consolidation_repair_failed request_id=%s unexpected_type=%s",
+                request_id,
+                type(response),
+            )
+            return None
+        return response
+
+    @staticmethod
+    def _find_suspicious_under_consumed_new_rows(
+        new_playbooks: list[UserPlaybook],
+        dedup_output: PlaybookConsolidationOutput,
+    ) -> list[str]:
+        candidates_by_id = {
+            f"NEW-{idx}": playbook for idx, playbook in enumerate(new_playbooks)
+        }
+        # Every kind except reject_new stores a row for its consumed candidates
+        # (raw for independent, refined-trigger copy for differentiate, the
+        # re-synthesized decision content for a sibling unify), so any of them
+        # can duplicate a unify's merged row. Only reject_new drops the row
+        # outright and is never a duplicate-storage risk worth a repair call.
+        consuming_kind_by_new_id = {
+            new_id: decision.kind
+            for decision in dedup_output.decisions
+            for new_id in _decision_new_ids(decision)
+        }
+        source_ids_by_decision = [
+            {
+                source_id
+                for new_id in _decision_new_ids(decision)
+                if new_id in candidates_by_id
+                for source_id in candidates_by_id[new_id].source_interaction_ids
+            }
+            for decision in dedup_output.decisions
+        ]
+        errors: list[str] = []
+
+        for decision_index, decision in enumerate(dedup_output.decisions):
+            if not isinstance(decision, UnifyDecision):
+                continue
+            consumed_ids = set(decision.new_ids)
+            consumed_source_ids = source_ids_by_decision[decision_index]
+            if not consumed_source_ids:
+                continue
+
+            # Candidates stored (near-)verbatim by another decision whose facts
+            # this unify's merged content may have absorbed.
+            for other_id, other in candidates_by_id.items():
+                if other_id in consumed_ids:
+                    continue
+                if consuming_kind_by_new_id.get(other_id) not in (
+                    "independent",
+                    "differentiate",
+                ):
+                    continue
+                if not consumed_source_ids.intersection(other.source_interaction_ids):
+                    continue
+                if check_string_token_overlap(decision.content, other.content):
+                    errors.append(
+                        f"decision[{decision_index}] unify likely absorbed facts from {other_id} but did not list it in new_id; "
+                        "if the unified rule absorbed this row's facts, add its id to new_id; "
+                        "if not, leave the decisions unchanged."
+                    )
+
+            # Sibling unify decisions each persist their own survivor row, so
+            # two same-source unifies with heavily overlapping FINAL contents
+            # are the same duplicate-storage class. Compare stored content to
+            # stored content and flag each pair once (j > i).
+            for sibling_index in range(decision_index + 1, len(dedup_output.decisions)):
+                sibling = dedup_output.decisions[sibling_index]
+                if not isinstance(sibling, UnifyDecision):
+                    continue
+                if not consumed_source_ids.intersection(
+                    source_ids_by_decision[sibling_index]
+                ):
+                    continue
+                if check_string_token_overlap(decision.content, sibling.content):
+                    errors.append(
+                        f"decision[{decision_index}] and decision[{sibling_index}] are both unify decisions over "
+                        "same-source NEW rows with highly overlapping content; if they describe the same skill, "
+                        "merge them into one unify decision whose new_id lists every consumed NEW id; "
+                        "if they are genuinely distinct skills, leave the decisions unchanged."
+                    )
+
+        return errors
+
     # ===============================
     # Apply path: discriminated-union decisions -> (new rows, archive ids)
     # ===============================
@@ -751,8 +1048,10 @@ class PlaybookConsolidator(BaseDeduplicator):
         new_playbooks: list[UserPlaybook],
         existing_playbooks: list[UserPlaybook],
         dedup_output: PlaybookConsolidationOutput,
-        request_id: str,
-        agent_version: str,  # noqa: ARG002
+        generation_request_id: str | None = None,
+        agent_version: str | None = None,  # noqa: ARG002
+        *,
+        request_id: str | None = None,
     ) -> tuple[list[UserPlaybook], list[int], list[tuple[int, list[int]]]]:
         """
         Build the deduplicated entry list from LLM decisions.
@@ -766,7 +1065,7 @@ class PlaybookConsolidator(BaseDeduplicator):
             new_playbooks: Flattened list of new (candidate) entries.
             existing_playbooks: List of existing entries from the DB.
             dedup_output: LLM decisions output (discriminated union).
-            request_id: Request ID stamped onto newly-built rows.
+            generation_request_id: Request ID stamped onto newly-built rows.
             agent_version: Agent version (currently unused, kept for symmetry).
 
         Returns:
@@ -781,6 +1080,10 @@ class PlaybookConsolidator(BaseDeduplicator):
             two rows (no single survivor), so its archived id appears in the
             delete set but NOT in any merge group.
         """
+        generation_request_id = _normalize_generation_request_id(
+            generation_request_id, request_id=request_id
+        )
+
         candidates_by_id = {
             f"NEW-{idx}": playbook for idx, playbook in enumerate(new_playbooks)
         }
@@ -810,7 +1113,7 @@ class PlaybookConsolidator(BaseDeduplicator):
                     existing_by_position=existing_by_position,
                     archive_ids=archive_ids,
                     seen_archive=seen_archive,
-                    request_id=request_id,
+                    generation_request_id=generation_request_id,
                 )
             except Exception as exc:  # noqa: BLE001 — per-decision isolation
                 result_counters.failed_count += 1
@@ -834,7 +1137,7 @@ class PlaybookConsolidator(BaseDeduplicator):
                     subsystem="playbook_consolidator",
                     op="apply_decision",
                     org_id=self.request_context.org_id,
-                    request_id=request_id,
+                    request_id=generation_request_id,
                     kind=decision.kind,
                     new_id=new_id_str,
                     existing_id=existing_id_str,
@@ -890,7 +1193,7 @@ class PlaybookConsolidator(BaseDeduplicator):
         existing_by_position: dict[str, UserPlaybook],
         archive_ids: list[int],
         seen_archive: set[int],
-        request_id: str,
+        generation_request_id: str,
     ) -> tuple[list[UserPlaybook], list[str], list[int]]:
         """Dispatch a single decision to its kind-specific apply method.
 
@@ -904,7 +1207,7 @@ class PlaybookConsolidator(BaseDeduplicator):
             archive_ids: Accumulator list mutated with ids to archive/delete.
             seen_archive: Accumulator set guarding ``archive_ids`` against
                 duplicate ids.
-            request_id: Request ID stamped onto newly-built rows.
+            generation_request_id: Request ID stamped onto newly-built rows.
 
         Returns:
             Tuple of ``(rows_to_insert, handled_new_ids, merge_source_ids)``.
@@ -923,7 +1226,7 @@ class PlaybookConsolidator(BaseDeduplicator):
                 existing_by_position=existing_by_position,
                 archive_ids=archive_ids,
                 seen_archive=seen_archive,
-                request_id=request_id,
+                generation_request_id=generation_request_id,
             )
         if isinstance(decision, RejectNewDecision):
             return self._apply_reject_new(
@@ -940,10 +1243,14 @@ class PlaybookConsolidator(BaseDeduplicator):
                 existing_by_position=existing_by_position,
                 archive_ids=archive_ids,
                 seen_archive=seen_archive,
-                request_id=request_id,
+                generation_request_id=generation_request_id,
             )
         if isinstance(decision, IndependentDecision):
-            return self._apply_independent(decision, candidates_by_id=candidates_by_id)
+            return self._apply_independent(
+                decision,
+                candidates_by_id=candidates_by_id,
+                generation_request_id=generation_request_id,
+            )
         raise ValueError(f"unknown decision kind: {decision}")
 
     def _apply_unify(
@@ -954,7 +1261,7 @@ class PlaybookConsolidator(BaseDeduplicator):
         existing_by_position: dict[str, UserPlaybook],
         archive_ids: list[int],
         seen_archive: set[int],
-        request_id: str,
+        generation_request_id: str,
     ) -> tuple[list[UserPlaybook], list[str], list[int]]:
         """Collapse / compose NEW (+ 0..N EXISTING) into one row.
 
@@ -974,7 +1281,7 @@ class PlaybookConsolidator(BaseDeduplicator):
             existing_by_position: Mapping ``"EXISTING-M"`` -> existing playbook.
             archive_ids: Accumulator mutated with EXISTING ids to archive.
             seen_archive: Dedup set for ``archive_ids``.
-            request_id: Request ID stamped on the unified row.
+            generation_request_id: Request ID stamped on the unified row.
 
         Returns:
             Tuple of ``([unified_row], [consumed NEW-N ids], merge_source_ids)``
@@ -1035,7 +1342,8 @@ class PlaybookConsolidator(BaseDeduplicator):
             user_playbook_id=0,
             user_id=primary_candidate.user_id,
             agent_version=primary_candidate.agent_version,
-            request_id=request_id,
+            # Legacy storage/API field; value may be synthetic for manual/rerun flows.
+            request_id=generation_request_id,
             playbook_name=primary_candidate.playbook_name,
             created_at=int(datetime.now(UTC).timestamp()),
             content=decision.content,
@@ -1120,7 +1428,7 @@ class PlaybookConsolidator(BaseDeduplicator):
         existing_by_position: dict[str, UserPlaybook],
         archive_ids: list[int],
         seen_archive: set[int],
-        request_id: str,
+        generation_request_id: str,
     ) -> tuple[list[UserPlaybook], list[str], list[int]]:
         """Archive the existing row and emit two refined rows in its place.
 
@@ -1136,7 +1444,7 @@ class PlaybookConsolidator(BaseDeduplicator):
             existing_by_position: Mapping ``"EXISTING-M"`` -> existing playbook.
             archive_ids: Accumulator mutated with the existing id to archive.
             seen_archive: Dedup set for ``archive_ids``.
-            request_id: Request ID stamped on both new rows.
+            generation_request_id: Request ID stamped on both new rows.
 
         Returns:
             Tuple of ``([refined_new_row, refined_existing_row], [NEW-N id],
@@ -1168,7 +1476,7 @@ class PlaybookConsolidator(BaseDeduplicator):
         refined_candidate = candidate.model_copy(
             update={
                 "user_playbook_id": 0,
-                "request_id": request_id,
+                "request_id": generation_request_id,
                 "trigger": decision.refined_new_trigger,
                 "created_at": now_ts,
             }
@@ -1176,7 +1484,7 @@ class PlaybookConsolidator(BaseDeduplicator):
         refined_existing = existing.model_copy(
             update={
                 "user_playbook_id": 0,
-                "request_id": request_id,
+                "request_id": generation_request_id,
                 "trigger": decision.refined_existing_trigger,
                 "created_at": now_ts,
                 "source_interaction_ids": list(existing.source_interaction_ids),
@@ -1189,6 +1497,7 @@ class PlaybookConsolidator(BaseDeduplicator):
         decision: IndependentDecision,
         *,
         candidates_by_id: dict[str, UserPlaybook],
+        generation_request_id: str,
     ) -> tuple[list[UserPlaybook], list[str], list[int]]:
         """Insert the new candidate unchanged; no archive.
 
@@ -1205,7 +1514,9 @@ class PlaybookConsolidator(BaseDeduplicator):
             candidate = candidates_by_id.get(new_id)
             if candidate is None:
                 raise KeyError(f"independent references unknown NEW id: {new_id}")
-            rows.append(candidate)
+            rows.append(
+                candidate.model_copy(update={"request_id": generation_request_id})
+            )
         return rows, decision.new_ids, []
 
     @staticmethod