From 262e6526b04551ae2668095114eb7c1ff31cf7de Mon Sep 17 00:00:00 2001 From: JAYARAM RADHAKRISHNAN Date: Sun, 15 Feb 2026 16:01:36 -0500 Subject: [PATCH 1/9] feat: persist task description in tip entity metadata generate_tips() now returns a TipGenerationResult containing both the tips and the source task_description. Both callers (PhoenixSync and MCP save_trajectory) store task_description in tip entity metadata, enabling future clustering of tips by task similarity. Trajectories without a task description default to "Task description unknown". --- kaizen/frontend/mcp/mcp_server.py | 5 +++-- kaizen/llm/tips/tips.py | 12 +++++++----- kaizen/schema/tips.py | 9 +++++++++ kaizen/sync/phoenix_sync.py | 9 +++++---- tests/unit/test_phoenix_sync.py | 7 ++++--- 5 files changed, 28 insertions(+), 14 deletions(-) diff --git a/kaizen/frontend/mcp/mcp_server.py b/kaizen/frontend/mcp/mcp_server.py index 62fcfdd4..4db9c0aa 100644 --- a/kaizen/frontend/mcp/mcp_server.py +++ b/kaizen/frontend/mcp/mcp_server.py @@ -98,7 +98,7 @@ def save_trajectory(trajectory_data: str, task_id: str | None = None) -> list[Re entities=entities, enable_conflict_resolution=False, ) - tips = generate_tips(messages) + result = generate_tips(messages) get_client().update_entities( namespace_id=kaizen_config.namespace_id, @@ -110,9 +110,10 @@ def save_trajectory(trajectory_data: str, task_id: str | None = None) -> list[Re "category": tip.category, "rationale": tip.rationale, "trigger": tip.trigger, + "task_description": result.task_description, }, ) - for tip in tips + for tip in result.tips ], enable_conflict_resolution=True, ) diff --git a/kaizen/llm/tips/tips.py b/kaizen/llm/tips/tips.py index 5335f79c..5d3afa40 100644 --- a/kaizen/llm/tips/tips.py +++ b/kaizen/llm/tips/tips.py @@ -8,7 +8,7 @@ from kaizen.config.llm import llm_settings from kaizen.utils.utils import clean_llm_response from kaizen.schema.exceptions import KaizenException -from kaizen.schema.tips import TipGenerationResponse, Tip +from kaizen.schema.tips import TipGenerationResponse, TipGenerationResult from pathlib import Path @@ -91,14 +91,14 @@ def parse_openai_agents_trajectory(messages: list[dict]) -> dict: steps_text.append(f"**Step {i} - Observation:**\n{content}") return { - "task_instruction": task_instruction or "Unknown task", + "task_instruction": task_instruction or "Task description unknown", "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"]]), } -def generate_tips(messages: list[dict]) -> list[Tip]: +def generate_tips(messages: list[dict]) -> TipGenerationResult: prompt_file = Path(__file__).parent / "prompts/generate_tips.jinja2" supported_params = get_supported_openai_params( model=llm_settings.tips_model, @@ -111,8 +111,9 @@ def generate_tips(messages: list[dict]) -> list[Tip]: ) constrained_decoding_supported = supports_response_format and response_schema_enabled trajectory_data = parse_openai_agents_trajectory(messages) + task_description = trajectory_data["task_instruction"] prompt = Template(prompt_file.read_text()).render( - task_instruction=trajectory_data["task_instruction"], + task_instruction=task_description, num_steps=trajectory_data["num_steps"], trajectory_summary=trajectory_data["trajectory_summary"], constrained_decoding_supported=constrained_decoding_supported, @@ -142,4 +143,5 @@ def generate_tips(messages: list[dict]) -> list[Tip]: .message.content ) clean_response = clean_llm_response(response) - return TipGenerationResponse.model_validate(json.loads(clean_response)).tips + tips = TipGenerationResponse.model_validate(json.loads(clean_response)).tips + return TipGenerationResult(tips=tips, task_description=task_description) diff --git a/kaizen/schema/tips.py b/kaizen/schema/tips.py index dabdf84d..ce9ac5e5 100644 --- a/kaizen/schema/tips.py +++ b/kaizen/schema/tips.py @@ -1,3 +1,4 @@ +from dataclasses import dataclass from pydantic import BaseModel, Field from typing import Literal @@ -11,3 +12,11 @@ class Tip(BaseModel): class TipGenerationResponse(BaseModel): tips: list[Tip] + + +@dataclass +class TipGenerationResult: + """Internal result from generate_tips(), pairing tips with the source task description.""" + + tips: list[Tip] + task_description: str diff --git a/kaizen/sync/phoenix_sync.py b/kaizen/sync/phoenix_sync.py index 991473f6..966bfa3e 100644 --- a/kaizen/sync/phoenix_sync.py +++ b/kaizen/sync/phoenix_sync.py @@ -469,9 +469,9 @@ def _process_trajectory(self, trajectory: dict) -> int: ) # Generate tips from the trajectory - tips = generate_tips(trajectory["messages"]) + result = generate_tips(trajectory["messages"]) - if tips: + if result.tips: tip_entities = [ Entity( type="guideline", @@ -482,9 +482,10 @@ def _process_trajectory(self, trajectory: dict) -> int: "trigger": tip.trigger, "source_trace_id": trajectory["trace_id"], "source_span_id": trajectory["span_id"], + "task_description": result.task_description, }, ) - for tip in tips + for tip in result.tips ] self.client.update_entities( namespace_id=self.namespace_id, @@ -492,7 +493,7 @@ def _process_trajectory(self, trajectory: dict) -> int: enable_conflict_resolution=True, ) - return len(tips) + return len(result.tips) def sync( self, diff --git a/tests/unit/test_phoenix_sync.py b/tests/unit/test_phoenix_sync.py index 695d7022..4d4fef00 100644 --- a/tests/unit/test_phoenix_sync.py +++ b/tests/unit/test_phoenix_sync.py @@ -6,6 +6,7 @@ import pytest from kaizen.sync.phoenix_sync import PhoenixSync, SyncResult +from kaizen.schema.tips import TipGenerationResult # Mark all tests in this module as phoenix tests (skipped by default) pytestmark = pytest.mark.phoenix @@ -507,7 +508,7 @@ def test_sync_includes_error_spans_when_requested(self, mock_generate_tips, mock mock_urlopen.return_value = mock_response phoenix_sync.client.search_entities.return_value = [] - mock_generate_tips.return_value = [] + mock_generate_tips.return_value = TipGenerationResult(tips=[], task_description="Task description unknown") result = phoenix_sync.sync(limit=10, include_errors=True) @@ -569,7 +570,7 @@ def test_sync_processes_valid_spans(self, mock_generate_tips, mock_urlopen, phoe mock_tip2.category = "optimization" mock_tip2.rationale = "Tip 2 rationale" mock_tip2.trigger = "Tip 2 trigger" - mock_generate_tips.return_value = [mock_tip1, mock_tip2] + mock_generate_tips.return_value = TipGenerationResult(tips=[mock_tip1, mock_tip2], task_description="Hello") result = phoenix_sync.sync(limit=10) @@ -623,7 +624,7 @@ def test_sync_returns_correct_counts(self, mock_generate_tips, mock_urlopen, pho mock_tip.category = "strategy" mock_tip.rationale = "Tip rationale" mock_tip.trigger = "Tip trigger" - mock_generate_tips.return_value = [mock_tip] + mock_generate_tips.return_value = TipGenerationResult(tips=[mock_tip], task_description="New message") result = phoenix_sync.sync(limit=10) From 5d6f9c2a0495dc3cfe4af2e0cab35d96d21f0a99 Mon Sep 17 00:00:00 2001 From: JAYARAM RADHAKRISHNAN Date: Sun, 15 Feb 2026 16:44:30 -0500 Subject: [PATCH 2/9] fix: guard against empty tips list in MCP save_trajectory Skip update_entities call when no tips are generated, aligning with the existing guard in phoenix_sync.py. --- kaizen/frontend/mcp/mcp_server.py | 35 ++++++++++++++++--------------- 1 file changed, 18 insertions(+), 17 deletions(-) diff --git a/kaizen/frontend/mcp/mcp_server.py b/kaizen/frontend/mcp/mcp_server.py index 4db9c0aa..850cc84f 100644 --- a/kaizen/frontend/mcp/mcp_server.py +++ b/kaizen/frontend/mcp/mcp_server.py @@ -100,23 +100,24 @@ def save_trajectory(trajectory_data: str, task_id: str | None = None) -> list[Re ) result = generate_tips(messages) - get_client().update_entities( - namespace_id=kaizen_config.namespace_id, - entities=[ - Entity( - type="guideline", - content=tip.content, - metadata={ - "category": tip.category, - "rationale": tip.rationale, - "trigger": tip.trigger, - "task_description": result.task_description, - }, - ) - for tip in result.tips - ], - enable_conflict_resolution=True, - ) + if result.tips: + get_client().update_entities( + namespace_id=kaizen_config.namespace_id, + entities=[ + Entity( + type="guideline", + content=tip.content, + metadata={ + "category": tip.category, + "rationale": tip.rationale, + "trigger": tip.trigger, + "task_description": result.task_description, + }, + ) + for tip in result.tips + ], + enable_conflict_resolution=True, + ) return get_client().search_entities( namespace_id=kaizen_config.namespace_id, From 626a0722d901afe919ca2e779cfdbcd51779c9fb Mon Sep 17 00:00:00 2001 From: JAYARAM RADHAKRISHNAN Date: Sun, 15 Feb 2026 16:45:49 -0500 Subject: [PATCH 3/9] test: assert task_description is persisted in tip entity metadata --- tests/unit/test_phoenix_sync.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/unit/test_phoenix_sync.py b/tests/unit/test_phoenix_sync.py index 4d4fef00..604ef568 100644 --- a/tests/unit/test_phoenix_sync.py +++ b/tests/unit/test_phoenix_sync.py @@ -578,6 +578,11 @@ def test_sync_processes_valid_spans(self, mock_generate_tips, mock_urlopen, phoe assert result.tips_generated == 2 phoenix_sync.client.update_entities.assert_called() + # Verify task_description is persisted in tip entity metadata + tip_update_call = phoenix_sync.client.update_entities.call_args_list[-1] + tip_entities = tip_update_call.kwargs["entities"] + assert all(e.metadata.get("task_description") == "Hello" for e in tip_entities) + @patch("kaizen.sync.phoenix_sync.urllib.request.urlopen") @patch("kaizen.sync.phoenix_sync.generate_tips") def test_sync_returns_correct_counts(self, mock_generate_tips, mock_urlopen, phoenix_sync): From fbbfdf6798c61f6f91bff0f564c9e5fb94858fb3 Mon Sep 17 00:00:00 2001 From: JAYARAM RADHAKRISHNAN Date: Sun, 15 Feb 2026 16:52:21 -0500 Subject: [PATCH 4/9] test: add unit tests for parse_openai_agents_trajectory fallback --- tests/unit/test_tips.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 tests/unit/test_tips.py diff --git a/tests/unit/test_tips.py b/tests/unit/test_tips.py new file mode 100644 index 00000000..152188a9 --- /dev/null +++ b/tests/unit/test_tips.py @@ -0,0 +1,25 @@ +"""Tests for tip generation utilities.""" + +import pytest + +from kaizen.llm.tips.tips import parse_openai_agents_trajectory + + +@pytest.mark.unit +class TestParseOpenaiAgentsTrajectory: + def test_extracts_task_instruction_from_first_user_message(self): + messages = [ + {"role": "user", "content": "Fix the login bug"}, + {"role": "assistant", "content": "I'll look into that."}, + ] + result = parse_openai_agents_trajectory(messages) + assert result["task_instruction"] == "Fix the login bug" + + def test_fallback_when_no_user_message(self): + messages = [{"role": "assistant", "content": "some response"}] + result = parse_openai_agents_trajectory(messages) + assert result["task_instruction"] == "Task description unknown" + + def test_fallback_when_empty_messages(self): + result = parse_openai_agents_trajectory([]) + assert result["task_instruction"] == "Task description unknown" From 76e9c834e3b2367b8f43b8c3f4e7e05fd81c919c Mon Sep 17 00:00:00 2001 From: JAYARAM RADHAKRISHNAN Date: Tue, 17 Feb 2026 10:55:34 -0500 Subject: [PATCH 5/9] feat: cluster tips by task description cosine similarity Add clustering module that embeds task descriptions via SentenceTransformer, computes pairwise cosine similarity, and groups tips using union-find. - New `cluster_entities()` function in `kaizen/llm/tips/clustering.py` - `KAIZEN_CLUSTERING_THRESHOLD` config (default 0.80) - `KaizenClient.cluster_tips()` method - `kaizen entities consolidate` CLI command (dry-run only for now) - 11 unit tests for clustering logic and union-find --- kaizen/cli/cli.py | 52 ++++++++ kaizen/config/kaizen.py | 1 + kaizen/frontend/client/kaizen_client.py | 20 +++ kaizen/llm/tips/clustering.py | 97 +++++++++++++++ tests/unit/test_clustering.py | 159 ++++++++++++++++++++++++ 5 files changed, 329 insertions(+) create mode 100644 kaizen/llm/tips/clustering.py create mode 100644 tests/unit/test_clustering.py diff --git a/kaizen/cli/cli.py b/kaizen/cli/cli.py index 2dfc68f0..035e1cdc 100644 --- a/kaizen/cli/cli.py +++ b/kaizen/cli/cli.py @@ -330,6 +330,58 @@ 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]") + + # ============================================================================= # 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..271feea6 100644 --- a/kaizen/frontend/client/kaizen_client.py +++ b/kaizen/frontend/client/kaizen_client.py @@ -70,6 +70,26 @@ 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 + ) -> 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. + + 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=10000) + return cluster_entities(entities, threshold=threshold) + # 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..a7b1ffff --- /dev/null +++ b/kaizen/llm/tips/clustering.py @@ -0,0 +1,97 @@ +"""Cluster tip entities by task description similarity.""" + +from __future__ import annotations + +import numpy as np +from sentence_transformers import SentenceTransformer + +from kaizen.schema.core import RecordedEntity + + +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 [] + + descriptions = [e.metadata["task_description"] for _, e in filtered] + + model = SentenceTransformer(embedding_model) + embeddings = model.encode(descriptions, normalize_embeddings=True) + similarity_matrix = np.asarray(embeddings) @ np.asarray(embeddings).T + + # Find pairs exceeding threshold + n = len(filtered) + pairs: list[tuple[int, int]] = [] + for i in range(n): + for j in range(i + 1, n): + if similarity_matrix[i, j] > threshold: + pairs.append((i, j)) + + 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 diff --git a/tests/unit/test_clustering.py b/tests/unit/test_clustering.py new file mode 100644 index 00000000..0618b05e --- /dev/null +++ b/tests/unit/test_clustering.py @@ -0,0 +1,159 @@ +"""Unit tests for tip clustering logic.""" + +from datetime import datetime +from unittest.mock import MagicMock, patch + +import numpy as np + +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 +# --------------------------------------------------------------------------- + + +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) + + +@patch("kaizen.llm.tips.clustering.SentenceTransformer") +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() + + def test_uses_default_embedding_model(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"), + _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() + call_arg = mock_st_cls.call_args[0][0] + assert "MiniLM" in call_arg or len(call_arg) > 0 From b54a52f3672acb9ff88c347d3e61cd72a915f9c4 Mon Sep 17 00:00:00 2001 From: JAYARAM RADHAKRISHNAN Date: Wed, 18 Feb 2026 08:30:21 -0500 Subject: [PATCH 6/9] feat: combine tips within clusters via LLM consolidation Add combine_cluster() that calls an LLM to merge overlapping tips in each cluster into fewer consolidated guidelines, with a 3-retry loop. Wire up the consolidate_tips() client method and CLI --no-dry-run path to delete originals and insert the combined results. --- kaizen/cli/cli.py | 15 ++ kaizen/frontend/client/kaizen_client.py | 51 +++++ kaizen/llm/tips/clustering.py | 97 ++++++++++ kaizen/llm/tips/prompts/combine_tips.jinja2 | 45 +++++ kaizen/schema/tips.py | 9 + tests/unit/test_combine_tips.py | 201 ++++++++++++++++++++ 6 files changed, 418 insertions(+) create mode 100644 kaizen/llm/tips/prompts/combine_tips.jinja2 create mode 100644 tests/unit/test_combine_tips.py diff --git a/kaizen/cli/cli.py b/kaizen/cli/cli.py index 035e1cdc..80d6a325 100644 --- a/kaizen/cli/cli.py +++ b/kaizen/cli/cli.py @@ -381,6 +381,21 @@ def consolidate_entities( 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/frontend/client/kaizen_client.py b/kaizen/frontend/client/kaizen_client.py index 271feea6..cf482a2d 100644 --- a/kaizen/frontend/client/kaizen_client.py +++ b/kaizen/frontend/client/kaizen_client.py @@ -1,6 +1,7 @@ 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 @@ -90,6 +91,56 @@ def cluster_tips( entities = self.get_all_entities(namespace_id, filters={"type": "guideline"}, limit=10000) 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) + tips_before = sum(len(c) for c in clusters) + tips_after = 0 + + for cluster in clusters: + consolidated_tips = combine_cluster(cluster) + tips_after += len(consolidated_tips) + + # Delete original entities + for entity in cluster: + self.delete_entity_by_id(namespace_id, entity.id) + + # Insert consolidated entities, preserving task_description from first entity + 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 new_entities: + self.update_entities(namespace_id, new_entities, enable_conflict_resolution=False) + + return ConsolidationResult( + clusters_found=len(clusters), + 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 index a7b1ffff..b56cd2d0 100644 --- a/kaizen/llm/tips/clustering.py +++ b/kaizen/llm/tips/clustering.py @@ -2,10 +2,20 @@ from __future__ import annotations +import json +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 def _union_find(n: int, pairs: list[tuple[int, int]]) -> list[list[int]]: @@ -95,3 +105,90 @@ def cluster_entities( 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. + """ + prompt_file = Path(__file__).parent / "prompts/combine_tips.jinja2" + + supported_params = get_supported_openai_params( + model=llm_settings.tips_model, + custom_llm_provider=llm_settings.custom_llm_provider, + ) + supports_response_format = 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 = Template(prompt_file.read_text()).render( + task_descriptions=task_descriptions, + tips=tips, + constrained_decoding_supported=constrained_decoding_supported, + ) + + last_error: Exception | None = None + for attempt in range(3): + try: + if constrained_decoding_supported: + litellm.enable_json_schema_validation = True + clean_response = ( + 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 + ) + else: + litellm.enable_json_schema_validation = False + response = ( + completion( + model=llm_settings.tips_model, + messages=[{"role": "user", "content": prompt}], + custom_llm_provider=llm_settings.custom_llm_provider, + ) + .choices[0] + .message.content + ) + clean_response = clean_llm_response(response) + + 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/schema/tips.py b/kaizen/schema/tips.py index ce9ac5e5..1d937a1e 100644 --- a/kaizen/schema/tips.py +++ b/kaizen/schema/tips.py @@ -20,3 +20,12 @@ class TipGenerationResult: tips: list[Tip] task_description: str + + +@dataclass +class ConsolidationResult: + """Summary of a tip consolidation run.""" + + clusters_found: int + tips_before: int + tips_after: int diff --git a/tests/unit/test_combine_tips.py b/tests/unit/test_combine_tips.py new file mode 100644 index 00000000..2af81b7d --- /dev/null +++ b/tests/unit/test_combine_tips.py @@ -0,0 +1,201 @@ +"""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 +# --------------------------------------------------------------------------- + + +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 +# --------------------------------------------------------------------------- + + +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 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 was called with consolidated entities + 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 + + @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 From 988a5ac8d6acf1e5e700e71ca088e9d5086f754d Mon Sep 17 00:00:00 2001 From: JAYARAM RADHAKRISHNAN Date: Thu, 19 Feb 2026 11:20:37 -0500 Subject: [PATCH 7/9] fix: address CodeRabbit review feedback on tip clustering MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Cache SentenceTransformer via lru_cache to avoid re-loading on every call - Guard O(n²) clustering with MAX_CLUSTER_ENTITIES=5000 cap - Vectorize pair scan with np.triu + np.where instead of nested loops - Use >= for inclusive threshold comparison - Extract DEFAULT_TASK_DESCRIPTION constant for consistent fallback - Reorder consolidate_tips to insert before delete to prevent data loss - Add per-cluster error handling so failed clusters are skipped gracefully - Hoist litellm.enable_json_schema_validation above the retry loop - Add configurable limit param to cluster_tips - Make TipGenerationResult and ConsolidationResult frozen dataclasses - Add @pytest.mark.unit to all clustering/combining test classes - Strengthen test_uses_default_embedding_model assertion --- kaizen/frontend/client/kaizen_client.py | 71 +++++++++++++++---------- kaizen/llm/tips/clustering.py | 36 +++++++++---- kaizen/llm/tips/tips.py | 4 +- kaizen/schema/tips.py | 6 ++- tests/unit/test_clustering.py | 13 +++-- tests/unit/test_combine_tips.py | 14 ++--- 6 files changed, 93 insertions(+), 51 deletions(-) diff --git a/kaizen/frontend/client/kaizen_client.py b/kaizen/frontend/client/kaizen_client.py index cf482a2d..2da90e74 100644 --- a/kaizen/frontend/client/kaizen_client.py +++ b/kaizen/frontend/client/kaizen_client.py @@ -1,3 +1,5 @@ +import logging + from kaizen.schema.core import Entity, Namespace, RecordedEntity from kaizen.schema.exceptions import NamespaceNotFoundException from kaizen.schema.conflict_resolution import EntityUpdate @@ -5,6 +7,8 @@ from kaizen.config.kaizen import KaizenConfig from kaizen.backend.base import BaseEntityBackend +logger = logging.getLogger(__name__) + class KaizenClient: """Wrapper client around kaizen entity backends.""" @@ -72,13 +76,14 @@ def delete_entity_by_id(self, namespace_id: str, entity_id: str) -> None: self.backend.delete_entity_by_id(namespace_id, entity_id) def cluster_tips( - self, namespace_id: str, threshold: float | None = None + 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. @@ -88,7 +93,7 @@ def cluster_tips( if threshold is None: threshold = self.config.clustering_threshold - entities = self.get_all_entities(namespace_id, filters={"type": "guideline"}, limit=10000) + entities = self.get_all_entities(namespace_id, filters={"type": "guideline"}, limit=limit) return cluster_entities(entities, threshold=threshold) def consolidate_tips( @@ -106,37 +111,49 @@ def consolidate_tips( from kaizen.llm.tips.clustering import combine_cluster clusters = self.cluster_tips(namespace_id, threshold=threshold) - tips_before = sum(len(c) for c in clusters) + clusters_found = 0 + tips_before = 0 tips_after = 0 for cluster in clusters: - consolidated_tips = combine_cluster(cluster) - tips_after += len(consolidated_tips) - - # Delete original entities - for entity in cluster: - self.delete_entity_by_id(namespace_id, entity.id) - - # Insert consolidated entities, preserving task_description from first entity - 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, - }, + try: + consolidated_tips = combine_cluster(cluster) + + # Insert consolidated entities first, preserving task_description from first entity + 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 new_entities: + self.update_entities(namespace_id, new_entities, enable_conflict_resolution=False) + + # Only delete originals after successful insert + for entity in cluster: + self.delete_entity_by_id(namespace_id, entity.id) + + clusters_found += 1 + tips_before += len(cluster) + tips_after += len(consolidated_tips) + 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, ) - for tip in consolidated_tips - ] - if new_entities: - self.update_entities(namespace_id, new_entities, enable_conflict_resolution=False) return ConsolidationResult( - clusters_found=len(clusters), + clusters_found=clusters_found, tips_before=tips_before, tips_after=tips_after, ) diff --git a/kaizen/llm/tips/clustering.py b/kaizen/llm/tips/clustering.py index b56cd2d0..33230c5c 100644 --- a/kaizen/llm/tips/clustering.py +++ b/kaizen/llm/tips/clustering.py @@ -3,6 +3,8 @@ from __future__ import annotations import json +import logging +from functools import lru_cache from pathlib import Path import litellm @@ -17,6 +19,15 @@ from kaizen.schema.tips import Tip, TipGenerationResponse from kaizen.utils.utils import clean_llm_response +logger = logging.getLogger(__name__) + +MAX_CLUSTER_ENTITIES = 5000 + + +@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. @@ -81,19 +92,26 @@ def cluster_entities( 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 = SentenceTransformer(embedding_model) + model = _get_sentence_transformer(embedding_model) embeddings = model.encode(descriptions, normalize_embeddings=True) similarity_matrix = np.asarray(embeddings) @ np.asarray(embeddings).T - # Find pairs exceeding threshold + # Find pairs meeting threshold (vectorized upper-triangle extraction) n = len(filtered) - pairs: list[tuple[int, int]] = [] - for i in range(n): - for j in range(i + 1, n): - if similarity_matrix[i, j] > threshold: - pairs.append((i, j)) + 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) @@ -157,11 +175,12 @@ def combine_cluster(entities: list[RecordedEntity]) -> list[Tip]: 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: - litellm.enable_json_schema_validation = True clean_response = ( completion( model=llm_settings.tips_model, @@ -173,7 +192,6 @@ def combine_cluster(entities: list[RecordedEntity]) -> list[Tip]: .message.content ) else: - litellm.enable_json_schema_validation = False response = ( completion( model=llm_settings.tips_model, diff --git a/kaizen/llm/tips/tips.py b/kaizen/llm/tips/tips.py index 5d3afa40..dd427c6d 100644 --- a/kaizen/llm/tips/tips.py +++ b/kaizen/llm/tips/tips.py @@ -8,7 +8,7 @@ from kaizen.config.llm import llm_settings from kaizen.utils.utils import clean_llm_response from kaizen.schema.exceptions import KaizenException -from kaizen.schema.tips import TipGenerationResponse, TipGenerationResult +from kaizen.schema.tips import DEFAULT_TASK_DESCRIPTION, TipGenerationResponse, TipGenerationResult from pathlib import Path @@ -91,7 +91,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 1d937a1e..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,7 +16,7 @@ class TipGenerationResponse(BaseModel): tips: list[Tip] -@dataclass +@dataclass(frozen=True) class TipGenerationResult: """Internal result from generate_tips(), pairing tips with the source task description.""" @@ -22,7 +24,7 @@ class TipGenerationResult: task_description: str -@dataclass +@dataclass(frozen=True) class ConsolidationResult: """Summary of a tip consolidation run.""" diff --git a/tests/unit/test_clustering.py b/tests/unit/test_clustering.py index 0618b05e..9959f1f9 100644 --- a/tests/unit/test_clustering.py +++ b/tests/unit/test_clustering.py @@ -4,6 +4,7 @@ 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 @@ -28,6 +29,7 @@ def _make_entity(entity_id: str, task_description: str | None = None) -> Recorde # --------------------------------------------------------------------------- +@pytest.mark.unit class TestUnionFind: def test_no_pairs(self): groups = _union_find(3, []) @@ -76,7 +78,8 @@ def _mock_encode(descriptions, normalize_embeddings=True): return np.array(vectors) -@patch("kaizen.llm.tips.clustering.SentenceTransformer") +@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() @@ -141,7 +144,9 @@ def test_single_entity(self, mock_st_cls): assert clusters == [] mock_st_cls.assert_not_called() - def test_uses_default_embedding_model(self, mock_st_cls): + @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 @@ -154,6 +159,4 @@ def test_uses_default_embedding_model(self, mock_st_cls): cluster_entities(entities, threshold=0.9) # Should use the default model from config - mock_st_cls.assert_called_once() - call_arg = mock_st_cls.call_args[0][0] - assert "MiniLM" in call_arg or len(call_arg) > 0 + 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 index 2af81b7d..1e3cad6d 100644 --- a/tests/unit/test_combine_tips.py +++ b/tests/unit/test_combine_tips.py @@ -56,6 +56,7 @@ def _mock_completion_response(tips: list[dict]) -> MagicMock: # --------------------------------------------------------------------------- +@pytest.mark.unit class TestCombineCluster: @patch("kaizen.llm.tips.clustering.completion") @patch("kaizen.llm.tips.clustering.supports_response_schema", return_value=False) @@ -126,6 +127,7 @@ def test_combine_cluster_uses_structured_output(self, _mock_params, _mock_schema # --------------------------------------------------------------------------- +@pytest.mark.unit class TestConsolidateTips: @patch("kaizen.llm.tips.clustering.combine_cluster") def test_consolidate_tips_deletes_originals_and_inserts_new(self, mock_combine): @@ -153,12 +155,7 @@ def test_consolidate_tips_deletes_originals_and_inserts_new(self, mock_combine): with patch.object(client, "cluster_tips", return_value=[entities_cluster]): client.consolidate_tips("test-ns") - # 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 was called with consolidated entities + # Verify insert was called before deletes (insert-first for safety) 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] @@ -168,6 +165,11 @@ def test_consolidate_tips_deletes_originals_and_inserts_new(self, mock_combine): 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") + @patch("kaizen.llm.tips.clustering.combine_cluster") def test_consolidate_tips_returns_correct_counts(self, mock_combine): # Cluster 1: 3 entities -> 1 consolidated tip From 8b2936c9f7cd9bac82e55277799cadf83013e23f Mon Sep 17 00:00:00 2001 From: JAYARAM RADHAKRISHNAN Date: Thu, 19 Feb 2026 11:34:51 -0500 Subject: [PATCH 8/9] fix: address CodeRabbit review feedback (round 3) - Separate consolidate_tips into two guarded scopes: combine+insert failures skip the cluster, while delete failures are logged independently without rolling back the successful insert - Hoist Jinja2 template compilation to module-level constant _COMBINE_TIPS_TEMPLATE to avoid re-reading/compiling on every call - Assert insert-before-delete ordering in test via mock_calls index comparison instead of just checking call counts --- kaizen/frontend/client/kaizen_client.py | 26 ++++++++++++++++--------- kaizen/llm/tips/clustering.py | 8 +++++--- tests/unit/test_combine_tips.py | 8 +++++++- 3 files changed, 29 insertions(+), 13 deletions(-) diff --git a/kaizen/frontend/client/kaizen_client.py b/kaizen/frontend/client/kaizen_client.py index 2da90e74..f9ce0e01 100644 --- a/kaizen/frontend/client/kaizen_client.py +++ b/kaizen/frontend/client/kaizen_client.py @@ -116,10 +116,10 @@ def consolidate_tips( tips_after = 0 for cluster in clusters: + # Phase 1: combine + insert (skip cluster on failure) try: consolidated_tips = combine_cluster(cluster) - # Insert consolidated entities first, preserving task_description from first entity task_description = (cluster[0].metadata or {}).get("task_description", "") new_entities = [ Entity( @@ -136,14 +136,6 @@ def consolidate_tips( ] if new_entities: self.update_entities(namespace_id, new_entities, enable_conflict_resolution=False) - - # Only delete originals after successful insert - for entity in cluster: - self.delete_entity_by_id(namespace_id, entity.id) - - clusters_found += 1 - tips_before += len(cluster) - tips_after += len(consolidated_tips) except Exception: logger.warning( "Failed to consolidate cluster of %d entities (IDs: %s); skipping.", @@ -151,6 +143,22 @@ def consolidate_tips( [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, diff --git a/kaizen/llm/tips/clustering.py b/kaizen/llm/tips/clustering.py index 33230c5c..c044186c 100644 --- a/kaizen/llm/tips/clustering.py +++ b/kaizen/llm/tips/clustering.py @@ -23,6 +23,10 @@ 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: @@ -139,8 +143,6 @@ def combine_cluster(entities: list[RecordedEntity]) -> list[Tip]: Raises: KaizenException: If the LLM call fails after 3 attempts. """ - prompt_file = Path(__file__).parent / "prompts/combine_tips.jinja2" - supported_params = get_supported_openai_params( model=llm_settings.tips_model, custom_llm_provider=llm_settings.custom_llm_provider, @@ -169,7 +171,7 @@ def combine_cluster(entities: list[RecordedEntity]) -> list[Tip]: for e in entities ] - prompt = Template(prompt_file.read_text()).render( + prompt = _COMBINE_TIPS_TEMPLATE.render( task_descriptions=task_descriptions, tips=tips, constrained_decoding_supported=constrained_decoding_supported, diff --git a/tests/unit/test_combine_tips.py b/tests/unit/test_combine_tips.py index 1e3cad6d..5b6b3790 100644 --- a/tests/unit/test_combine_tips.py +++ b/tests/unit/test_combine_tips.py @@ -155,7 +155,7 @@ def test_consolidate_tips_deletes_originals_and_inserts_new(self, mock_combine): with patch.object(client, "cluster_tips", return_value=[entities_cluster]): client.consolidate_tips("test-ns") - # Verify insert was called before deletes (insert-first for safety) + # 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] @@ -170,6 +170,12 @@ def test_consolidate_tips_deletes_originals_and_inserts_new(self, mock_combine): 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 From fc2000ac963fd3692c215265e06d9de73aa0d474 Mon Sep 17 00:00:00 2001 From: JAYARAM RADHAKRISHNAN Date: Fri, 20 Feb 2026 13:24:53 -0500 Subject: [PATCH 9/9] fix: address CodeRabbit review feedback (round 4) - Warn when cluster_tips hits the entity fetch limit - Skip deletion when LLM returns no consolidated tips for a cluster - Guard against None .message.content in combine_cluster --- kaizen/frontend/client/kaizen_client.py | 23 +++++++++++++++-------- kaizen/llm/tips/clustering.py | 25 +++++++++++++------------ 2 files changed, 28 insertions(+), 20 deletions(-) diff --git a/kaizen/frontend/client/kaizen_client.py b/kaizen/frontend/client/kaizen_client.py index f9ce0e01..6de2fdb6 100644 --- a/kaizen/frontend/client/kaizen_client.py +++ b/kaizen/frontend/client/kaizen_client.py @@ -75,9 +75,7 @@ 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]]: + 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: @@ -94,11 +92,15 @@ def cluster_tips( 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: + def consolidate_tips(self, namespace_id: str, threshold: float | None = None) -> ConsolidationResult: """Cluster similar tips and combine each cluster into consolidated guidelines. Args: @@ -134,8 +136,13 @@ def consolidate_tips( ) for tip in consolidated_tips ] - if new_entities: - self.update_entities(namespace_id, new_entities, enable_conflict_resolution=False) + 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.", diff --git a/kaizen/llm/tips/clustering.py b/kaizen/llm/tips/clustering.py index c044186c..29ee7e41 100644 --- a/kaizen/llm/tips/clustering.py +++ b/kaizen/llm/tips/clustering.py @@ -23,9 +23,7 @@ MAX_CLUSTER_ENTITIES = 5000 -_COMBINE_TIPS_TEMPLATE = Template( - (Path(__file__).parent / "prompts/combine_tips.jinja2").read_text() -) +_COMBINE_TIPS_TEMPLATE = Template((Path(__file__).parent / "prompts/combine_tips.jinja2").read_text()) @lru_cache(maxsize=4) @@ -147,7 +145,7 @@ def combine_cluster(entities: list[RecordedEntity]) -> list[Tip]: model=llm_settings.tips_model, custom_llm_provider=llm_settings.custom_llm_provider, ) - supports_response_format = supported_params and "response_format" in supported_params + 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, @@ -155,11 +153,9 @@ def combine_cluster(entities: list[RecordedEntity]) -> list[Tip]: 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") - )) + task_descriptions = list( + dict.fromkeys((e.metadata or {}).get("task_description", "") for e in entities if (e.metadata or {}).get("task_description")) + ) tips = [ { @@ -183,7 +179,7 @@ def combine_cluster(entities: list[RecordedEntity]) -> list[Tip]: for attempt in range(3): try: if constrained_decoding_supported: - clean_response = ( + content = ( completion( model=llm_settings.tips_model, messages=[{"role": "user", "content": prompt}], @@ -193,8 +189,11 @@ def combine_cluster(entities: list[RecordedEntity]) -> list[Tip]: .choices[0] .message.content ) + if content is None: + raise KaizenException("LLM returned None content for combine_cluster") + clean_response = content else: - response = ( + content = ( completion( model=llm_settings.tips_model, messages=[{"role": "user", "content": prompt}], @@ -203,7 +202,9 @@ def combine_cluster(entities: list[RecordedEntity]) -> list[Tip]: .choices[0] .message.content ) - clean_response = clean_llm_response(response) + 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: