alvis
9ddeea6cac
Each unique task title is now enriched by LiteLLM once and cached in the DB. Subsequent agent compute cycles (every 12h) fetch the cache before calling ml-serving; only new titles hit the tip-generator. - DB: task_enrichments(content_hash PK, description, model, created_at) - TS: fetchEnrichmentCache / persistEnrichments helpers in agent-outputs.ts; enrichment_cache passed in compute request, new_enrichments persisted from response - Python: AgentComputeRequest.enrichment_cache / AgentComputeResponse.new_enrichments; AgentInput.enrichment_cache; _enrich_batch returns (descriptions, new_entries); cluster_tasks returns (clusters, new_enrichments) - FocusAreaAgent stashes new_enrichments in signals_snapshot under _new_enrichments; compute_agent endpoint pops it before storing the snapshot Closes part of #129 Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
210 lines
8.9 KiB
Python
210 lines
8.9 KiB
Python
"""Unit tests for ml.agents.clustering (issue #97, #129).
|
|
|
|
LLM and embedding calls are mocked so tests run without Ollama or LiteLLM.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import sys, os
|
|
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, _enrich_batch
|
|
|
|
|
|
# ── helpers ──────────────────────────────────────────────────────────────────
|
|
|
|
def _task(content: str, project_id: str | None = None, is_overdue: bool = False) -> dict:
|
|
t: dict = {"content": content, "is_overdue": is_overdue}
|
|
if project_id:
|
|
t["project_id"] = project_id
|
|
return t
|
|
|
|
|
|
def _embed_seq(*vecs):
|
|
"""Return a side_effect list so successive _embed calls return these vectors."""
|
|
return list(vecs)
|
|
|
|
|
|
# ── Cluster dataclass ─────────────────────────────────────────────────────────
|
|
|
|
class TestCluster:
|
|
def test_task_count(self):
|
|
c = Cluster(label="X", tasks=[_task("a"), _task("b")])
|
|
assert c.task_count == 2
|
|
|
|
def test_overdue_count(self):
|
|
c = Cluster(label="X", tasks=[_task("a", is_overdue=True), _task("b")])
|
|
assert c.overdue_count == 1
|
|
|
|
|
|
# ── cosine similarity ─────────────────────────────────────────────────────────
|
|
|
|
class TestCosine:
|
|
def test_identical_vectors(self):
|
|
v = [1.0, 0.0, 0.0]
|
|
assert _cosine(v, v) == 1.0
|
|
|
|
def test_orthogonal_vectors(self):
|
|
assert _cosine([1.0, 0.0], [0.0, 1.0]) == 0.0
|
|
|
|
def test_zero_vector(self):
|
|
assert _cosine([0.0, 0.0], [1.0, 0.0]) == 0.0
|
|
|
|
|
|
# ── greedy clustering ─────────────────────────────────────────────────────────
|
|
|
|
class TestGreedyClustering:
|
|
def _similar_vec(self, base: list[float], noise: float = 0.01) -> list[float]:
|
|
return [x + noise for x in base]
|
|
|
|
def test_similar_tasks_grouped(self):
|
|
v = [1.0, 0.0, 0.0]
|
|
v2 = [0.999, 0.001, 0.0]
|
|
items = [
|
|
(_task("A"), v),
|
|
(_task("B"), v2),
|
|
]
|
|
clusters = _greedy_cluster(items)
|
|
assert len(clusters) == 1
|
|
assert clusters[0].task_count == 2
|
|
|
|
def test_dissimilar_tasks_separate(self):
|
|
v1 = [1.0, 0.0, 0.0]
|
|
v2 = [0.0, 1.0, 0.0]
|
|
items = [(_task("A"), v1), (_task("B"), v2)]
|
|
clusters = _greedy_cluster(items)
|
|
assert len(clusters) == 2
|
|
|
|
def test_label_from_first_task(self):
|
|
v = [1.0, 0.0]
|
|
clusters = _greedy_cluster([(_task("Write report"), v)])
|
|
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, new = _enrich_batch(["Buy milk", "Fix bug"])
|
|
assert result == ["Buy milk", "Fix bug"] and new == {}
|
|
|
|
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, new = _enrich_batch(["Buy milk"])
|
|
assert result == ["Expanded description."]
|
|
assert len(new) == 1
|
|
|
|
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, new = _enrich_batch(["Buy milk"])
|
|
assert result == ["Buy milk"]
|
|
assert new == {} # failed enrichments are not persisted
|
|
|
|
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, new = _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"]
|
|
|
|
def test_uses_persistent_cache(self, monkeypatch):
|
|
monkeypatch.setenv("LITELLM_URL", "http://fake-litellm")
|
|
from ml.agents.clustering import _content_hash
|
|
h = _content_hash("Buy milk")
|
|
call_count = {"n": 0}
|
|
def fake_enrich(title, url):
|
|
call_count["n"] += 1
|
|
return "new desc"
|
|
with patch("ml.agents.clustering._enrich_title", side_effect=fake_enrich):
|
|
result, new = _enrich_batch(["Buy milk"], persistent_cache={h: "cached desc"})
|
|
assert call_count["n"] == 0 # cache hit, no LLM call
|
|
assert result == ["cached desc"]
|
|
assert new == {}
|
|
|
|
|
|
# ── cluster_tasks integration ─────────────────────────────────────────────────
|
|
|
|
class TestClusterTasks:
|
|
def _no_enrich(self, titles, persistent_cache=None):
|
|
return titles, {}
|
|
|
|
def test_empty_tasks(self):
|
|
clusters, new = cluster_tasks([])
|
|
assert clusters == [] and new == {}
|
|
|
|
def test_fallback_when_embed_unavailable(self):
|
|
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
|
|
labels = {c.label for c in clusters}
|
|
assert "p1" in labels and "p2" in labels
|
|
|
|
def test_fallback_groups_by_project(self):
|
|
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}
|
|
assert by_label["work"] == 3
|
|
assert by_label["home"] == 2
|
|
|
|
def test_tasks_without_content_go_to_other(self):
|
|
v = [1.0, 0.0]
|
|
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}
|
|
assert "Other tasks" in labels
|
|
|
|
def test_semantic_clustering_groups_similar(self):
|
|
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._enrich_batch", side_effect=self._no_enrich), \
|
|
patch("ml.agents.clustering._embed_batch", return_value=batch_result):
|
|
tasks = [
|
|
_task("Write report"),
|
|
_task("Review PR"),
|
|
_task("Buy groceries"),
|
|
_task("Cook dinner"),
|
|
]
|
|
clusters, _ = cluster_tasks(tasks)
|
|
assert len(clusters) == 2
|
|
assert all(c.task_count == 2 for c in clusters)
|
|
|
|
def test_all_tasks_no_content_fallback_by_project(self):
|
|
tasks = [{"project_id": "p1", "is_overdue": False},
|
|
{"project_id": "p2", "is_overdue": False}]
|
|
clusters, new = cluster_tasks(tasks)
|
|
assert len(clusters) == 2 and new == {}
|
|
|
|
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."]
|
|
|
|
def test_new_enrichments_returned(self):
|
|
v = [1.0, 0.0]
|
|
with patch("ml.agents.clustering._enrich_batch", return_value=(["desc"], {"abc123": "desc"})), \
|
|
patch("ml.agents.clustering._embed_batch", return_value=[v]):
|
|
_, new = cluster_tasks([_task("Buy milk")])
|
|
assert new == {"abc123": "desc"}
|