OpenJarvis/tests/learning/training/test_data.py

"""Tests for TrainingDataMiner — SFT, routing, and agent config extraction."""

from __future__ import annotations

import time
from typing import Any, List

from openjarvis.core.types import StepType, Trace, TraceStep
from openjarvis.learning.training.data import TrainingDataMiner

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_trace(
    *,
    query: str = "Hello world",
    agent: str = "simple",
    model: str = "qwen3:8b",
    engine: str = "ollama",
    result: str = "Hi there!",
    feedback: float | None = 0.9,
    outcome: str | None = "success",
    tools: List[str] | None = None,
) -> Trace:
    """Build a Trace with optional TOOL_CALL steps."""
    now = time.time()
    steps: list[TraceStep] = [
        TraceStep(
            step_type=StepType.GENERATE,
            timestamp=now,
            duration_seconds=0.5,
            input={"prompt": query},
            output={"text": result, "tokens": 10},
        ),
    ]
    if tools:
        for tool_name in tools:
            steps.append(
                TraceStep(
                    step_type=StepType.TOOL_CALL,
                    timestamp=now + 0.1,
                    duration_seconds=0.1,
                    input={"tool": tool_name, "args": {}},
                    output={"result": "ok"},
                )
            )
    steps.append(
        TraceStep(
            step_type=StepType.RESPOND,
            timestamp=now + 1.0,
            duration_seconds=0.0,
            input={},
            output={"text": result},
        )
    )
    return Trace(
        query=query,
        agent=agent,
        model=model,
        engine=engine,
        result=result,
        feedback=feedback,
        outcome=outcome,
        started_at=now,
        ended_at=now + 1.0,
        total_tokens=10,
        total_latency_seconds=1.0,
        steps=steps,
    )


class FakeTraceStore:
    """Minimal mock that satisfies TrainingDataMiner's needs."""

    def __init__(self, traces: list[Trace] | None = None):
        self._traces = traces or []

    def list_traces(self, *, limit: int = 10000, **kwargs: Any) -> list[Trace]:
        return self._traces[:limit]


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


class TestExtractSFTPairs:
    def test_extract_sft_pairs_from_successful_traces(self) -> None:
        """SFT pairs are extracted from high-quality traces."""
        traces = [
            _make_trace(
                query="Write a hello world in Python",
                result="print('hello')",
                feedback=0.9,
            ),
            _make_trace(query="Solve x^2=4", result="x=2 or x=-2", feedback=0.8),
        ]
        store = FakeTraceStore(traces)
        miner = TrainingDataMiner(store)
        pairs = miner.extract_sft_pairs()

        assert len(pairs) == 2
        # Check structure of first pair
        p0 = pairs[0]
        assert p0["input"] == "Write a hello world in Python"
        assert p0["output"] == "print('hello')"
        assert "query_class" in p0
        assert p0["model"] == "qwen3:8b"
        assert p0["feedback"] == 0.9

    def test_deduplication(self) -> None:
        """Duplicate (input, output) pairs are collapsed to a single entry."""
        traces = [
            _make_trace(query="Hi", result="Hello!", feedback=0.9),
            _make_trace(query="Hi", result="Hello!", feedback=0.95),
        ]
        store = FakeTraceStore(traces)
        miner = TrainingDataMiner(store)
        pairs = miner.extract_sft_pairs()

        assert len(pairs) == 1

    def test_min_quality_filter(self) -> None:
        """Traces below min_quality are excluded from SFT pairs."""
        traces = [
            _make_trace(query="Good", result="Fine", feedback=0.9),
            _make_trace(query="Bad", result="Nope", feedback=0.3),
            _make_trace(query="None", result="Null", feedback=None),
        ]
        store = FakeTraceStore(traces)
        miner = TrainingDataMiner(store, min_quality=0.7)
        pairs = miner.extract_sft_pairs()

        assert len(pairs) == 1
        assert pairs[0]["input"] == "Good"


class TestExtractRoutingPairs:
    def test_extract_routing_pairs(self) -> None:
        """Routing pairs group traces by query class and find best model."""
        traces = [
            _make_trace(query="def foo(): pass", model="codellama:7b", feedback=0.95),
            _make_trace(
                query="import os; print(os.getcwd())",
                model="codellama:7b",
                feedback=0.85,
            ),
            _make_trace(query="def bar(): return 1", model="qwen3:8b", feedback=0.7),
        ]
        store = FakeTraceStore(traces)
        miner = TrainingDataMiner(store, min_quality=0.7)
        routing = miner.extract_routing_pairs()

        assert "code" in routing
        code_entry = routing["code"]
        assert code_entry["best_model"] == "codellama:7b"
        assert code_entry["sample_count"] == 3
        assert "codellama:7b" in code_entry["all_models"]
        assert "qwen3:8b" in code_entry["all_models"]


class TestExtractAgentConfigPairs:
    def test_extract_agent_config_pairs(self) -> None:
        """Agent config pairs find best agent and tools per query class."""
        traces = [
            _make_trace(
                query="Calculate 2+2",
                agent="orchestrator",
                tools=["calculator"],
                feedback=0.95,
            ),
            _make_trace(
                query="Compute 3*3",
                agent="orchestrator",
                tools=["calculator", "think"],
                feedback=0.9,
            ),
            _make_trace(
                query="Solve x+1=3",
                agent="simple",
                feedback=0.6,
            ),
        ]
        store = FakeTraceStore(traces)
        miner = TrainingDataMiner(store, min_quality=0.5)
        agent_cfg = miner.extract_agent_config_pairs()

        assert "math" in agent_cfg
        math_entry = agent_cfg["math"]
        assert math_entry["best_agent"] == "orchestrator"
        assert "calculator" in math_entry["best_tools"]
        assert math_entry["sample_count"] == 3


class TestOutcomeFilter:
    def test_failure_traces_excluded_despite_high_feedback(self) -> None:
        """Traces with outcome='failure' are excluded even if feedback is high."""
        traces = [
            _make_trace(
                query="Good query",
                result="Good answer",
                feedback=0.9,
                outcome="success",
            ),
            _make_trace(
                query="Bad query",
                result="Bad answer",
                feedback=0.9,
                outcome="failure",
            ),
        ]
        store = FakeTraceStore(traces)
        miner = TrainingDataMiner(store, min_quality=0.7)

        sft = miner.extract_sft_pairs()
        assert len(sft) == 1
        assert sft[0]["input"] == "Good query"

        routing = miner.extract_routing_pairs()
        total = sum(v["sample_count"] for v in routing.values())
        assert total == 1

        agent_cfg = miner.extract_agent_config_pairs()
        total_agent = sum(v["sample_count"] for v in agent_cfg.values())
        assert total_agent == 1


class TestEmptyStore:
    def test_empty_store_returns_empty(self) -> None:
        """All extractors return empty results for an empty store."""
        store = FakeTraceStore([])
        miner = TrainingDataMiner(store)

        assert miner.extract_sft_pairs() == []
        assert miner.extract_routing_pairs() == {}
        assert miner.extract_agent_config_pairs() == {}