feat(clustering): LLM-enrichment before embedding (port from taskpile #129)
Ported from taskpile experiments/clustering_eval (prompt v1, qwen2.5:1.5b). The experiment showed ARI 0.22→0.77 and AUROC 0.76→0.91 on synthetic tasks when embedding LLM-expanded descriptions instead of raw titles. - Expand each task title via LiteLLM tip-generator before embedding - Prefix with "clustering: " (nomic-embed-text task instruction prefix) - Cache expansions in-memory by content hash within a compute cycle - Falls back to raw title if enrichment fails; no change to fallback behaviour Fixes #129 Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -1,6 +1,6 @@
|
||||
"""Unit tests for ml.agents.clustering (issue #97).
|
||||
"""Unit tests for ml.agents.clustering (issue #97, #129).
|
||||
|
||||
Embedding calls are mocked so tests run without Ollama.
|
||||
LLM and embedding calls are mocked so tests run without Ollama or LiteLLM.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -9,7 +9,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", ".."))
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from ml.agents.clustering import cluster_tasks, Cluster, _greedy_cluster, _cosine, _embed_batch
|
||||
from ml.agents.clustering import cluster_tasks, Cluster, _greedy_cluster, _cosine, _embed_batch, _enrich_batch
|
||||
|
||||
|
||||
# ── helpers ──────────────────────────────────────────────────────────────────
|
||||
@@ -82,15 +82,51 @@ class TestGreedyClustering:
|
||||
assert clusters[0].label == "Write report"
|
||||
|
||||
|
||||
# ── enrichment ───────────────────────────────────────────────────────────────
|
||||
|
||||
class TestEnrichBatch:
|
||||
def test_falls_back_to_raw_when_no_litellm_url(self, monkeypatch):
|
||||
monkeypatch.delenv("LITELLM_URL", raising=False)
|
||||
result = _enrich_batch(["Buy milk", "Fix bug"])
|
||||
assert result == ["Buy milk", "Fix bug"]
|
||||
|
||||
def test_uses_description_when_litellm_available(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_URL", "http://fake-litellm")
|
||||
with patch("ml.agents.clustering._enrich_title", return_value="Expanded description."):
|
||||
result = _enrich_batch(["Buy milk"])
|
||||
assert result == ["Expanded description."]
|
||||
|
||||
def test_falls_back_to_raw_title_on_enrich_failure(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_URL", "http://fake-litellm")
|
||||
with patch("ml.agents.clustering._enrich_title", return_value=None):
|
||||
result = _enrich_batch(["Buy milk"])
|
||||
assert result == ["Buy milk"]
|
||||
|
||||
def test_deduplicates_identical_titles(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_URL", "http://fake-litellm")
|
||||
call_count = {"n": 0}
|
||||
def fake_enrich(title, url):
|
||||
call_count["n"] += 1
|
||||
return f"desc:{title}"
|
||||
with patch("ml.agents.clustering._enrich_title", side_effect=fake_enrich):
|
||||
result = _enrich_batch(["Buy milk", "Buy milk", "Fix bug"])
|
||||
assert call_count["n"] == 2 # only 2 unique titles
|
||||
assert result == ["desc:Buy milk", "desc:Buy milk", "desc:Fix bug"]
|
||||
|
||||
|
||||
# ── cluster_tasks integration ─────────────────────────────────────────────────
|
||||
|
||||
class TestClusterTasks:
|
||||
def _no_enrich(self, titles):
|
||||
return titles # pass-through; enrichment tested separately
|
||||
|
||||
def test_empty_tasks(self):
|
||||
result = cluster_tasks([])
|
||||
assert result == []
|
||||
|
||||
def test_fallback_when_embed_unavailable(self):
|
||||
with patch("ml.agents.clustering._embed_batch", return_value=None):
|
||||
with patch("ml.agents.clustering._enrich_batch", side_effect=self._no_enrich), \
|
||||
patch("ml.agents.clustering._embed_batch", return_value=None):
|
||||
tasks = [_task("A", "p1"), _task("B", "p2"), _task("C", "p1")]
|
||||
clusters = cluster_tasks(tasks)
|
||||
assert len(clusters) == 2
|
||||
@@ -98,7 +134,8 @@ class TestClusterTasks:
|
||||
assert "p1" in labels and "p2" in labels
|
||||
|
||||
def test_fallback_groups_by_project(self):
|
||||
with patch("ml.agents.clustering._embed_batch", return_value=None):
|
||||
with patch("ml.agents.clustering._enrich_batch", side_effect=self._no_enrich), \
|
||||
patch("ml.agents.clustering._embed_batch", return_value=None):
|
||||
tasks = [_task("A", "work")] * 3 + [_task("B", "home")] * 2
|
||||
clusters = cluster_tasks(tasks)
|
||||
by_label = {c.label: c.task_count for c in clusters}
|
||||
@@ -107,7 +144,8 @@ class TestClusterTasks:
|
||||
|
||||
def test_tasks_without_content_go_to_other(self):
|
||||
v = [1.0, 0.0]
|
||||
with patch("ml.agents.clustering._embed_batch", return_value=[v]):
|
||||
with patch("ml.agents.clustering._enrich_batch", side_effect=self._no_enrich), \
|
||||
patch("ml.agents.clustering._embed_batch", return_value=[v]):
|
||||
tasks = [_task("Has content"), {"is_overdue": False}]
|
||||
clusters = cluster_tasks(tasks)
|
||||
labels = {c.label for c in clusters}
|
||||
@@ -117,7 +155,8 @@ class TestClusterTasks:
|
||||
v_work = [1.0, 0.0, 0.0]
|
||||
v_home = [0.0, 1.0, 0.0]
|
||||
batch_result = [v_work, v_work, v_home, v_home]
|
||||
with patch("ml.agents.clustering._embed_batch", return_value=batch_result):
|
||||
with patch("ml.agents.clustering._enrich_batch", side_effect=self._no_enrich), \
|
||||
patch("ml.agents.clustering._embed_batch", return_value=batch_result):
|
||||
tasks = [
|
||||
_task("Write report"),
|
||||
_task("Review PR"),
|
||||
@@ -133,3 +172,15 @@ class TestClusterTasks:
|
||||
{"project_id": "p2", "is_overdue": False}]
|
||||
clusters = cluster_tasks(tasks)
|
||||
assert len(clusters) == 2
|
||||
|
||||
def test_enrich_called_before_embed(self):
|
||||
"""Verify enrichment output (not raw title) is what gets embedded."""
|
||||
v = [1.0, 0.0]
|
||||
captured = {}
|
||||
def fake_embed(texts):
|
||||
captured["texts"] = texts
|
||||
return [v] * len(texts)
|
||||
with patch("ml.agents.clustering._enrich_batch", return_value=["Expanded desc."]), \
|
||||
patch("ml.agents.clustering._embed_batch", side_effect=fake_embed):
|
||||
cluster_tasks([_task("Buy milk")])
|
||||
assert captured["texts"] == ["clustering: Expanded desc."]
|
||||
|
||||
Reference in New Issue
Block a user