oO/ml/agents/tests/test_clustering.py

"""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 = _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._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 = 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."]