diff --git a/kaizen/cli/cli.py b/kaizen/cli/cli.py index 2dfc68f0..80d6a325 100644 --- a/kaizen/cli/cli.py +++ b/kaizen/cli/cli.py @@ -330,6 +330,73 @@ def show_entity( raise typer.Exit(1) +@entities_app.command("consolidate") +def consolidate_entities( + namespace: Annotated[str, typer.Argument(help="Namespace to consolidate entities in")], + threshold: Annotated[Optional[float], typer.Option("--threshold", "-t", help="Cosine similarity threshold (0-1)")] = None, + dry_run: Annotated[bool, typer.Option("--dry-run", help="Show clusters without modifying anything")] = True, +): + """Cluster similar guideline entities by task description similarity.""" + from kaizen.config.kaizen import kaizen_config + + client = get_client() + + effective_threshold = threshold if threshold is not None else kaizen_config.clustering_threshold + + console.print(f"[bold]Clustering entities in '{namespace}'[/bold]") + console.print(f" Threshold: {effective_threshold}") + console.print(f" Dry run: {dry_run}") + console.print() + + try: + clusters = client.cluster_tips(namespace, threshold=effective_threshold) + except NamespaceNotFoundException: + console.print(f"[red]Namespace '{namespace}' not found.[/red]") + raise typer.Exit(1) + + if not clusters: + console.print("[yellow]No clusters found. Tips have dissimilar task descriptions.[/yellow]") + return + + console.print(f"[green]Found {len(clusters)} cluster(s)[/green]\n") + + for i, cluster in enumerate(clusters, 1): + table = Table(title=f"Cluster {i} ({len(cluster)} entities)") + table.add_column("ID", style="cyan", max_width=20) + table.add_column("Task Description", max_width=40) + table.add_column("Content", max_width=50) + + for entity in cluster: + task_desc = (entity.metadata or {}).get("task_description", "") + if len(task_desc) > 40: + task_desc = task_desc[:37] + "..." + content_str = str(entity.content) + if len(content_str) > 50: + content_str = content_str[:47] + "..." + table.add_row(str(entity.id), task_desc, content_str) + + console.print(table) + console.print() + + total_entities = sum(len(c) for c in clusters) + console.print(f"[dim]Total: {total_entities} entities in {len(clusters)} clusters[/dim]") + + if dry_run: + console.print("\n[yellow]Dry run — no changes made. Use --no-dry-run to consolidate.[/yellow]") + return + + console.print("\n[bold]Consolidating clusters...[/bold]") + try: + result = client.consolidate_tips(namespace, threshold=effective_threshold) + console.print("[green]Consolidation complete:[/green]") + console.print(f" Clusters combined: {result.clusters_found}") + console.print(f" Tips before: {result.tips_before}") + console.print(f" Tips after: {result.tips_after}") + except KaizenException as e: + console.print(f"[red]Consolidation failed: {e}[/red]") + raise typer.Exit(1) + + # ============================================================================= # Sync Commands # ============================================================================= diff --git a/kaizen/config/kaizen.py b/kaizen/config/kaizen.py index 265a3054..bcb53542 100644 --- a/kaizen/config/kaizen.py +++ b/kaizen/config/kaizen.py @@ -7,6 +7,7 @@ class KaizenConfig(BaseSettings): backend: Literal["milvus", "filesystem"] = "milvus" namespace_id: str = "kaizen" settings: BaseSettings | None = None + clustering_threshold: float = 0.80 # to reload settings call kaizen_config.__init__() diff --git a/kaizen/frontend/client/kaizen_client.py b/kaizen/frontend/client/kaizen_client.py index 7da08f43..6de2fdb6 100644 --- a/kaizen/frontend/client/kaizen_client.py +++ b/kaizen/frontend/client/kaizen_client.py @@ -1,9 +1,14 @@ +import logging + from kaizen.schema.core import Entity, Namespace, RecordedEntity from kaizen.schema.exceptions import NamespaceNotFoundException from kaizen.schema.conflict_resolution import EntityUpdate +from kaizen.schema.tips import ConsolidationResult from kaizen.config.kaizen import KaizenConfig from kaizen.backend.base import BaseEntityBackend +logger = logging.getLogger(__name__) + class KaizenClient: """Wrapper client around kaizen entity backends.""" @@ -70,6 +75,104 @@ def delete_entity_by_id(self, namespace_id: str, entity_id: str) -> None: """Delete a specific entity by its ID.""" self.backend.delete_entity_by_id(namespace_id, entity_id) + def cluster_tips(self, namespace_id: str, threshold: float | None = None, limit: int = 10000) -> list[list[RecordedEntity]]: + """Cluster guideline entities by task description similarity. + + Args: + namespace_id: Namespace to fetch entities from. + threshold: Cosine similarity threshold (0-1). Defaults to config value. + limit: Maximum number of guideline entities to fetch for clustering. + + Returns: + List of clusters, each containing related RecordedEntity objects. + """ + from kaizen.llm.tips.clustering import cluster_entities + + if threshold is None: + threshold = self.config.clustering_threshold + + entities = self.get_all_entities(namespace_id, filters={"type": "guideline"}, limit=limit) + if len(entities) >= limit: + logger.warning( + "Fetched %d entities (hit limit=%d); clustering results may be incomplete. Consider increasing the limit.", + len(entities), + limit, + ) + return cluster_entities(entities, threshold=threshold) + + def consolidate_tips(self, namespace_id: str, threshold: float | None = None) -> ConsolidationResult: + """Cluster similar tips and combine each cluster into consolidated guidelines. + + Args: + namespace_id: Namespace to consolidate entities in. + threshold: Cosine similarity threshold (0-1). Defaults to config value. + + Returns: + ConsolidationResult with counts of clusters, tips before, and tips after. + """ + from kaizen.llm.tips.clustering import combine_cluster + + clusters = self.cluster_tips(namespace_id, threshold=threshold) + clusters_found = 0 + tips_before = 0 + tips_after = 0 + + for cluster in clusters: + # Phase 1: combine + insert (skip cluster on failure) + try: + consolidated_tips = combine_cluster(cluster) + + task_description = (cluster[0].metadata or {}).get("task_description", "") + new_entities = [ + Entity( + content=tip.content, + type="guideline", + metadata={ + "task_description": task_description, + "rationale": tip.rationale, + "category": tip.category, + "trigger": tip.trigger, + }, + ) + for tip in consolidated_tips + ] + if not new_entities: + logger.warning( + "LLM returned no consolidated tips for cluster (IDs: %s); skipping deletion.", + [e.id for e in cluster], + ) + continue + self.update_entities(namespace_id, new_entities, enable_conflict_resolution=False) + except Exception: + logger.warning( + "Failed to consolidate cluster of %d entities (IDs: %s); skipping.", + len(cluster), + [e.id for e in cluster], + exc_info=True, + ) + continue + + clusters_found += 1 + tips_before += len(cluster) + tips_after += len(consolidated_tips) + + # Phase 2: delete originals (log errors but don't roll back insert) + for entity in cluster: + try: + self.delete_entity_by_id(namespace_id, entity.id) + except Exception: + logger.warning( + "Failed to delete original entity %s after successful insert; skipping.", + entity.id, + exc_info=True, + ) + + return ConsolidationResult( + clusters_found=clusters_found, + tips_before=tips_before, + tips_after=tips_after, + ) + # Convenience methods for common patterns def namespace_exists(self, namespace_id: str) -> bool: """Check if a namespace exists.""" diff --git a/kaizen/llm/tips/clustering.py b/kaizen/llm/tips/clustering.py new file mode 100644 index 00000000..29ee7e41 --- /dev/null +++ b/kaizen/llm/tips/clustering.py @@ -0,0 +1,215 @@ +"""Cluster tip entities by task description similarity.""" + +from __future__ import annotations + +import json +import logging +from functools import lru_cache +from pathlib import Path + +import litellm +import numpy as np +from jinja2 import Template +from litellm import completion, get_supported_openai_params, supports_response_schema +from sentence_transformers import SentenceTransformer + +from kaizen.config.llm import llm_settings +from kaizen.schema.core import RecordedEntity +from kaizen.schema.exceptions import KaizenException +from kaizen.schema.tips import Tip, TipGenerationResponse +from kaizen.utils.utils import clean_llm_response + +logger = logging.getLogger(__name__) + +MAX_CLUSTER_ENTITIES = 5000 + +_COMBINE_TIPS_TEMPLATE = Template((Path(__file__).parent / "prompts/combine_tips.jinja2").read_text()) + + +@lru_cache(maxsize=4) +def _get_sentence_transformer(model_name: str) -> SentenceTransformer: + return SentenceTransformer(model_name) + + +def _union_find(n: int, pairs: list[tuple[int, int]]) -> list[list[int]]: + """Group indices into connected components using union-find with path compression. + + Args: + n: Total number of elements. + pairs: Index pairs (i, j) to union together. + + Returns: + List of groups, where each group is a list of indices. + """ + parent = list(range(n)) + + def find(x: int) -> int: + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + for i, j in pairs: + ri, rj = find(i), find(j) + if ri != rj: + parent[ri] = rj + + groups: dict[int, list[int]] = {} + for i in range(n): + root = find(i) + groups.setdefault(root, []).append(i) + + return list(groups.values()) + + +def cluster_entities( + entities: list[RecordedEntity], + threshold: float = 0.80, + embedding_model: str | None = None, +) -> list[list[RecordedEntity]]: + """Cluster entities by cosine similarity of their task descriptions. + + Args: + entities: Guideline entities with optional ``task_description`` in metadata. + threshold: Cosine similarity threshold for clustering (0-1). + embedding_model: SentenceTransformer model name. Defaults to the model + configured in ``kaizen.config.milvus``. + + Returns: + List of clusters (each a list of ``RecordedEntity``), excluding + single-entity clusters. + """ + if embedding_model is None: + from kaizen.config.milvus import milvus_other_settings + + embedding_model = milvus_other_settings.embedding_model + + # Filter to entities that have a task_description + filtered: list[tuple[int, RecordedEntity]] = [] + for idx, entity in enumerate(entities): + td = (entity.metadata or {}).get("task_description") + if td: + filtered.append((idx, entity)) + + if len(filtered) < 2: + return [] + + if len(filtered) > MAX_CLUSTER_ENTITIES: + logger.warning( + "Too many entities for clustering (%d > %d). Truncating to first %d.", + len(filtered), + MAX_CLUSTER_ENTITIES, + MAX_CLUSTER_ENTITIES, + ) + filtered = filtered[:MAX_CLUSTER_ENTITIES] + + descriptions = [e.metadata["task_description"] for _, e in filtered] + + model = _get_sentence_transformer(embedding_model) + embeddings = model.encode(descriptions, normalize_embeddings=True) + similarity_matrix = np.asarray(embeddings) @ np.asarray(embeddings).T + + # Find pairs meeting threshold (vectorized upper-triangle extraction) + n = len(filtered) + mask = np.triu(similarity_matrix >= threshold, k=1) + rows, cols = np.where(mask) + pairs: list[tuple[int, int]] = list(zip(rows.tolist(), cols.tolist())) + + groups = _union_find(n, pairs) + + # Convert index groups back to entity clusters, excluding singletons + clusters: list[list[RecordedEntity]] = [] + for group in groups: + if len(group) < 2: + continue + clusters.append([filtered[i][1] for i in group]) + + return clusters + + +def combine_cluster(entities: list[RecordedEntity]) -> list[Tip]: + """Combine tips from a cluster of related entities into consolidated guidelines. + + Uses an LLM to merge overlapping tips into fewer, non-redundant guidelines. + + Args: + entities: Cluster of related entities to combine. + + Returns: + Consolidated list of tips. + + Raises: + KaizenException: If the LLM call fails after 3 attempts. + """ + supported_params = get_supported_openai_params( + model=llm_settings.tips_model, + custom_llm_provider=llm_settings.custom_llm_provider, + ) + supports_response_format = bool(supported_params and "response_format" in supported_params) + response_schema_enabled = supports_response_schema( + model=llm_settings.tips_model, + custom_llm_provider=llm_settings.custom_llm_provider, + ) + constrained_decoding_supported = supports_response_format and response_schema_enabled + + # Deduplicate task descriptions + task_descriptions = list( + dict.fromkeys((e.metadata or {}).get("task_description", "") for e in entities if (e.metadata or {}).get("task_description")) + ) + + tips = [ + { + "content": str(e.content), + "rationale": (e.metadata or {}).get("rationale", ""), + "category": (e.metadata or {}).get("category", "strategy"), + "trigger": (e.metadata or {}).get("trigger", ""), + } + for e in entities + ] + + prompt = _COMBINE_TIPS_TEMPLATE.render( + task_descriptions=task_descriptions, + tips=tips, + constrained_decoding_supported=constrained_decoding_supported, + ) + + litellm.enable_json_schema_validation = constrained_decoding_supported + + last_error: Exception | None = None + for attempt in range(3): + try: + if constrained_decoding_supported: + content = ( + completion( + model=llm_settings.tips_model, + messages=[{"role": "user", "content": prompt}], + response_format=TipGenerationResponse, + custom_llm_provider=llm_settings.custom_llm_provider, + ) + .choices[0] + .message.content + ) + if content is None: + raise KaizenException("LLM returned None content for combine_cluster") + clean_response = content + else: + content = ( + completion( + model=llm_settings.tips_model, + messages=[{"role": "user", "content": prompt}], + custom_llm_provider=llm_settings.custom_llm_provider, + ) + .choices[0] + .message.content + ) + if content is None: + raise KaizenException("LLM returned None content for combine_cluster") + clean_response = clean_llm_response(content) + + return TipGenerationResponse.model_validate(json.loads(clean_response)).tips + except Exception as e: + last_error = e + if attempt < 2: + continue + + raise KaizenException("Failed to combine cluster tips after 3 attempts") from last_error diff --git a/kaizen/llm/tips/prompts/combine_tips.jinja2 b/kaizen/llm/tips/prompts/combine_tips.jinja2 new file mode 100644 index 00000000..a463a0e2 --- /dev/null +++ b/kaizen/llm/tips/prompts/combine_tips.jinja2 @@ -0,0 +1,45 @@ +You are consolidating multiple AI agent guidelines that were generated from similar tasks. + +# Task Descriptions +These guidelines came from tasks like: +{% for desc in task_descriptions %} +- {{ desc }} +{% endfor %} + +# Existing Guidelines +{% for tip in tips %} +## Guideline {{ loop.index }} +- **Content:** {{ tip.content }} +- **Rationale:** {{ tip.rationale }} +- **Category:** {{ tip.category }} +- **Trigger:** {{ tip.trigger }} + +{% endfor %} + +# Your Task +Combine the above guidelines into a smaller set of HIGH-QUALITY, CONSOLIDATED, NON-REDUNDANT guidelines. Merge overlapping advice, remove duplicates, and preserve any unique insights. + +**Things to Remember:** +1. Each consolidated guideline should be self-contained, clear and actionable +2. Preserve important nuances from different source guidelines. Preserve specific technical details (API names, parameter names, etc.) +3. Use the most specific and helpful phrasing +4. Remove guidelines that are too generic or obvious. +5. Reduce the total count while retaining all valuable information + +{% if not constrained_decoding_supported %} +**Output Format (JSON):** +```json +{ + "tips": [ + { + "content": "Clear, actionable tip", + "rationale": "Why this tip helps", + "category": "strategy|recovery|optimization", + "trigger": "When to apply this tip" + } + ] +} +``` + +Generate consolidated guidelines now. Return ONLY the JSON, no other text. +{% endif %} diff --git a/kaizen/llm/tips/tips.py b/kaizen/llm/tips/tips.py index 88af0bd5..6c63a274 100644 --- a/kaizen/llm/tips/tips.py +++ b/kaizen/llm/tips/tips.py @@ -10,7 +10,7 @@ from kaizen.config.llm import llm_settings from kaizen.schema.exceptions import KaizenException -from kaizen.schema.tips import TipGenerationResponse, TipGenerationResult +from kaizen.schema.tips import DEFAULT_TASK_DESCRIPTION, TipGenerationResponse, TipGenerationResult from kaizen.utils.utils import clean_llm_response logger = logging.getLogger(__name__) @@ -96,7 +96,7 @@ def parse_openai_agents_trajectory(messages: list[dict]) -> dict: steps_text.append(f"**Step {i} - Observation:**\n{content}") return { - "task_instruction": task_instruction or "Task description unknown", + "task_instruction": task_instruction or DEFAULT_TASK_DESCRIPTION, "trajectory_summary": "\n\n".join(steps_text), "function_calls": function_calls, "num_steps": len([s for s in agent_steps if s["type"] in ["action", "reasoning"]]), diff --git a/kaizen/schema/tips.py b/kaizen/schema/tips.py index ce9ac5e5..34a6c9b2 100644 --- a/kaizen/schema/tips.py +++ b/kaizen/schema/tips.py @@ -2,6 +2,8 @@ from pydantic import BaseModel, Field from typing import Literal +DEFAULT_TASK_DESCRIPTION = "Task description unknown" + class Tip(BaseModel): content: str = Field(description="Clear, actionable tip") @@ -14,9 +16,18 @@ class TipGenerationResponse(BaseModel): tips: list[Tip] -@dataclass +@dataclass(frozen=True) class TipGenerationResult: """Internal result from generate_tips(), pairing tips with the source task description.""" tips: list[Tip] task_description: str + + +@dataclass(frozen=True) +class ConsolidationResult: + """Summary of a tip consolidation run.""" + + clusters_found: int + tips_before: int + tips_after: int diff --git a/tests/unit/test_clustering.py b/tests/unit/test_clustering.py new file mode 100644 index 00000000..9959f1f9 --- /dev/null +++ b/tests/unit/test_clustering.py @@ -0,0 +1,162 @@ +"""Unit tests for tip clustering logic.""" + +from datetime import datetime +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from kaizen.llm.tips.clustering import _union_find, cluster_entities +from kaizen.schema.core import RecordedEntity + + +def _make_entity(entity_id: str, task_description: str | None = None) -> RecordedEntity: + """Helper to create a RecordedEntity with optional task_description metadata.""" + metadata = {} + if task_description is not None: + metadata["task_description"] = task_description + return RecordedEntity( + id=entity_id, + content=f"Tip for {entity_id}", + type="guideline", + metadata=metadata, + created_at=datetime(2025, 1, 1), + ) + + +# --------------------------------------------------------------------------- +# _union_find tests +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestUnionFind: + def test_no_pairs(self): + groups = _union_find(3, []) + assert len(groups) == 3 + for g in groups: + assert len(g) == 1 + + def test_single_pair(self): + groups = _union_find(3, [(0, 1)]) + sizes = sorted(len(g) for g in groups) + assert sizes == [1, 2] + + def test_transitive_merge(self): + groups = _union_find(4, [(0, 1), (1, 2)]) + sizes = sorted(len(g) for g in groups) + assert sizes == [1, 3] + + def test_all_connected(self): + groups = _union_find(3, [(0, 1), (1, 2)]) + sizes = sorted(len(g) for g in groups) + assert sizes == [3] + + def test_two_components(self): + groups = _union_find(4, [(0, 1), (2, 3)]) + sizes = sorted(len(g) for g in groups) + assert sizes == [2, 2] + + +# --------------------------------------------------------------------------- +# cluster_entities tests +# --------------------------------------------------------------------------- + + +def _mock_encode(descriptions, normalize_embeddings=True): + """Return controlled embeddings: identical vectors for similar, orthogonal for different.""" + vectors = [] + for desc in descriptions: + if "error handling" in desc.lower(): + vectors.append([1.0, 0.0, 0.0]) + elif "caching" in desc.lower(): + vectors.append([0.0, 1.0, 0.0]) + elif "logging" in desc.lower(): + vectors.append([0.0, 0.0, 1.0]) + else: + vectors.append([0.5, 0.5, 0.0]) + return np.array(vectors) + + +@pytest.mark.unit +@patch("kaizen.llm.tips.clustering._get_sentence_transformer") +class TestClusterEntities: + def test_groups_similar_tasks(self, mock_st_cls): + mock_model = MagicMock() + mock_model.encode = _mock_encode + mock_st_cls.return_value = mock_model + + entities = [ + _make_entity("1", "Improve error handling in API"), + _make_entity("2", "Better error handling for edge cases"), + _make_entity("3", "Add caching to database queries"), + ] + + clusters = cluster_entities(entities, threshold=0.9, embedding_model="test-model") + + # The two error handling entities should cluster together + assert len(clusters) == 1 + cluster_ids = {e.id for e in clusters[0]} + assert cluster_ids == {"1", "2"} + + def test_separates_different_tasks(self, mock_st_cls): + mock_model = MagicMock() + mock_model.encode = _mock_encode + mock_st_cls.return_value = mock_model + + entities = [ + _make_entity("1", "Improve error handling in API"), + _make_entity("2", "Add caching to database queries"), + _make_entity("3", "Set up logging infrastructure"), + ] + + clusters = cluster_entities(entities, threshold=0.9, embedding_model="test-model") + + # All orthogonal — no clusters + assert clusters == [] + + def test_skips_missing_task_description(self, mock_st_cls): + mock_model = MagicMock() + mock_model.encode = _mock_encode + mock_st_cls.return_value = mock_model + + entities = [ + _make_entity("1", "Improve error handling in API"), + _make_entity("2", None), # no task_description + _make_entity("3", "Better error handling for edge cases"), + ] + + clusters = cluster_entities(entities, threshold=0.9, embedding_model="test-model") + + # Entity 2 is excluded; entities 1 and 3 should cluster + assert len(clusters) == 1 + cluster_ids = {e.id for e in clusters[0]} + assert cluster_ids == {"1", "3"} + + def test_empty_input(self, mock_st_cls): + clusters = cluster_entities([], threshold=0.8, embedding_model="test-model") + assert clusters == [] + mock_st_cls.assert_not_called() + + def test_single_entity(self, mock_st_cls): + entities = [_make_entity("1", "Some task")] + clusters = cluster_entities(entities, threshold=0.8, embedding_model="test-model") + assert clusters == [] + mock_st_cls.assert_not_called() + + @patch("kaizen.config.milvus.milvus_other_settings") + def test_uses_default_embedding_model(self, mock_settings, mock_st_cls): + mock_settings.embedding_model = "test-default-model" + mock_model = MagicMock() + mock_model.encode = _mock_encode + mock_st_cls.return_value = mock_model + + entities = [ + _make_entity("1", "Improve error handling"), + _make_entity("2", "Better error handling"), + ] + + cluster_entities(entities, threshold=0.9) + + # Should use the default model from config + mock_st_cls.assert_called_once_with("test-default-model") diff --git a/tests/unit/test_combine_tips.py b/tests/unit/test_combine_tips.py new file mode 100644 index 00000000..5b6b3790 --- /dev/null +++ b/tests/unit/test_combine_tips.py @@ -0,0 +1,209 @@ +"""Unit tests for tip combining and consolidation logic.""" + +import json +from datetime import datetime +from unittest.mock import MagicMock, patch + +import pytest + +from kaizen.llm.tips.clustering import combine_cluster +from kaizen.schema.core import RecordedEntity +from kaizen.schema.exceptions import KaizenException +from kaizen.schema.tips import Tip, ConsolidationResult + + +def _make_entity(entity_id: str, content: str, task_description: str = "do a task") -> RecordedEntity: + return RecordedEntity( + id=entity_id, + content=content, + type="guideline", + metadata={ + "task_description": task_description, + "rationale": "some rationale", + "category": "strategy", + "trigger": "when needed", + }, + created_at=datetime(2025, 1, 1), + ) + + +def _mock_completion_response(tips: list[dict]) -> MagicMock: + """Build a mock litellm completion response.""" + response = MagicMock() + response.choices = [MagicMock()] + response.choices[0].message.content = json.dumps({"tips": tips}) + return response + + +SAMPLE_TIPS = [ + { + "content": "Use retry logic for flaky APIs", + "rationale": "APIs can fail transiently", + "category": "recovery", + "trigger": "When calling external APIs", + }, + { + "content": "Log errors with context", + "rationale": "Easier debugging", + "category": "optimization", + "trigger": "When handling exceptions", + }, +] + + +# --------------------------------------------------------------------------- +# combine_cluster tests +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestCombineCluster: + @patch("kaizen.llm.tips.clustering.completion") + @patch("kaizen.llm.tips.clustering.supports_response_schema", return_value=False) + @patch("kaizen.llm.tips.clustering.get_supported_openai_params", return_value=[]) + def test_combine_cluster_returns_tips(self, _mock_params, _mock_schema, mock_completion): + mock_completion.return_value = _mock_completion_response(SAMPLE_TIPS) + + entities = [ + _make_entity("1", "Always retry on failure"), + _make_entity("2", "Add error logging"), + ] + + result = combine_cluster(entities) + + assert len(result) == 2 + assert all(isinstance(t, Tip) for t in result) + assert result[0].content == "Use retry logic for flaky APIs" + assert result[1].category == "optimization" + mock_completion.assert_called_once() + + @patch("kaizen.llm.tips.clustering.completion") + @patch("kaizen.llm.tips.clustering.supports_response_schema", return_value=False) + @patch("kaizen.llm.tips.clustering.get_supported_openai_params", return_value=[]) + def test_combine_cluster_retries_on_failure(self, _mock_params, _mock_schema, mock_completion): + mock_completion.side_effect = [ + ValueError("bad json"), + ValueError("bad json again"), + _mock_completion_response(SAMPLE_TIPS[:1]), + ] + + entities = [_make_entity("1", "Tip A"), _make_entity("2", "Tip B")] + result = combine_cluster(entities) + + assert len(result) == 1 + assert result[0].content == "Use retry logic for flaky APIs" + assert mock_completion.call_count == 3 + + @patch("kaizen.llm.tips.clustering.completion") + @patch("kaizen.llm.tips.clustering.supports_response_schema", return_value=False) + @patch("kaizen.llm.tips.clustering.get_supported_openai_params", return_value=[]) + def test_combine_cluster_raises_after_max_retries(self, _mock_params, _mock_schema, mock_completion): + mock_completion.side_effect = ValueError("always fails") + + entities = [_make_entity("1", "Tip A"), _make_entity("2", "Tip B")] + + with pytest.raises(KaizenException, match="Failed to combine cluster tips after 3 attempts"): + combine_cluster(entities) + + assert mock_completion.call_count == 3 + + @patch("kaizen.llm.tips.clustering.completion") + @patch("kaizen.llm.tips.clustering.supports_response_schema", return_value=True) + @patch("kaizen.llm.tips.clustering.get_supported_openai_params", return_value=["response_format"]) + def test_combine_cluster_uses_structured_output(self, _mock_params, _mock_schema, mock_completion): + mock_completion.return_value = _mock_completion_response(SAMPLE_TIPS[:1]) + + entities = [_make_entity("1", "Tip A"), _make_entity("2", "Tip B")] + result = combine_cluster(entities) + + assert len(result) == 1 + # Verify response_format was passed + _, kwargs = mock_completion.call_args + assert "response_format" in kwargs + + +# --------------------------------------------------------------------------- +# consolidate_tips tests +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +class TestConsolidateTips: + @patch("kaizen.llm.tips.clustering.combine_cluster") + def test_consolidate_tips_deletes_originals_and_inserts_new(self, mock_combine): + consolidated = [ + Tip(content="Combined tip", rationale="Merged", category="strategy", trigger="Always"), + ] + mock_combine.return_value = consolidated + + entities_cluster = [ + _make_entity("1", "Tip A", "error handling"), + _make_entity("2", "Tip B", "error handling"), + ] + + mock_backend = MagicMock() + mock_backend.search_entities.return_value = entities_cluster + + from kaizen.frontend.client.kaizen_client import KaizenClient + + client = KaizenClient.__new__(KaizenClient) + client.backend = mock_backend + client.config = MagicMock() + client.config.clustering_threshold = 0.80 + + # Mock cluster_tips to return our cluster + with patch.object(client, "cluster_tips", return_value=[entities_cluster]): + client.consolidate_tips("test-ns") + + # Verify insert was called with correct args + assert mock_backend.update_entities.call_count == 1 + call_args = mock_backend.update_entities.call_args + ns_id, new_entities, enable_cr = call_args[0] + assert ns_id == "test-ns" + assert len(new_entities) == 1 + assert new_entities[0].content == "Combined tip" + assert new_entities[0].metadata["task_description"] == "error handling" + assert enable_cr is False + + # Verify deletes were called for each original entity + assert mock_backend.delete_entity_by_id.call_count == 2 + mock_backend.delete_entity_by_id.assert_any_call("test-ns", "1") + mock_backend.delete_entity_by_id.assert_any_call("test-ns", "2") + + # Verify insert happened before deletes + call_names = [str(c) for c in mock_backend.mock_calls] + insert_idx = next(i for i, c in enumerate(call_names) if "update_entities" in c) + first_delete_idx = next(i for i, c in enumerate(call_names) if "delete_entity_by_id" in c) + assert insert_idx < first_delete_idx + + @patch("kaizen.llm.tips.clustering.combine_cluster") + def test_consolidate_tips_returns_correct_counts(self, mock_combine): + # Cluster 1: 3 entities -> 1 consolidated tip + # Cluster 2: 2 entities -> 2 consolidated tips + mock_combine.side_effect = [ + [Tip(content="C1", rationale="R", category="strategy", trigger="T")], + [ + Tip(content="C2a", rationale="R", category="strategy", trigger="T"), + Tip(content="C2b", rationale="R", category="optimization", trigger="T"), + ], + ] + + cluster1 = [_make_entity(f"c1-{i}", f"Tip {i}", "task A") for i in range(3)] + cluster2 = [_make_entity(f"c2-{i}", f"Tip {i}", "task B") for i in range(2)] + + mock_backend = MagicMock() + + from kaizen.frontend.client.kaizen_client import KaizenClient + + client = KaizenClient.__new__(KaizenClient) + client.backend = mock_backend + client.config = MagicMock() + client.config.clustering_threshold = 0.80 + + with patch.object(client, "cluster_tips", return_value=[cluster1, cluster2]): + result = client.consolidate_tips("test-ns") + + assert isinstance(result, ConsolidationResult) + assert result.clusters_found == 2 + assert result.tips_before == 5 + assert result.tips_after == 3