agents/tests/test_intent_agent.py

"""
IntentAgent unit tests — covering gaps not addressed by test_hitl.py
(clarification routing) or test_model_routing.py (Haiku/Sonnet routing).

Covers:
  • Malformed / non-JSON LLM response → AgentResult(success=False)
  • Partial JSON (missing required fields) → AgentResult(success=False)
  • context.task_spec is populated on success
  • NER intent → task_type="token_classification", label_names extracted
  • Result message contains task_type on success
  • base_model_hint appears in result output
  • System prompt is ≥ 1 024 tokens (cache eligibility)
"""
from __future__ import annotations

import json
from unittest.mock import AsyncMock, MagicMock

import pytest

from agents.base import AgentContext


# ── Helpers ───────────────────────────────────────────────────────────────────

def _fake_response(text: str) -> MagicMock:
    usage = MagicMock()
    usage.input_tokens = 300
    usage.output_tokens = 60
    usage.cache_read_input_tokens = 0
    usage.cache_creation_input_tokens = 0
    content = MagicMock()
    content.text = text
    resp = MagicMock()
    resp.usage = usage
    resp.content = [content]
    return resp


def _valid_json(
    task_type: str = "text_classification",
    num_labels: int = 3,
    label_names: list[str] | None = None,
    input_col: str = "text",
    label_col: str = "label",
    base_model: str = "distilbert-base-uncased",
    confidence: float = 0.92,
    clarification: str | None = None,
) -> str:
    return json.dumps({
        "task_type": task_type,
        "num_labels": num_labels,
        "label_names": label_names or ["a", "b", "c"],
        "input_column": input_col,
        "label_column": label_col,
        "base_model_hint": base_model,
        "confidence": confidence,
        "clarification_needed": clarification,
    })


def _make_agent(response_text: str):
    mock_client = AsyncMock()
    mock_client.messages.create = AsyncMock(return_value=_fake_response(response_text))
    from agents.intent import IntentAgent
    return IntentAgent(client=mock_client)


# ── JSON parse failures ───────────────────────────────────────────────────────

class TestIntentAgentParseFailures:
    @pytest.mark.asyncio
    async def test_plain_text_response_returns_failure(self):
        """LLM returns prose instead of JSON → AgentResult(success=False)."""
        agent = _make_agent("Sure! I'll classify your emails.")
        ctx = AgentContext(run_id="r1", user_intent="classify emails")
        result = await agent.run(ctx)

        assert result.success is False
        assert ctx.task_spec == {} or ctx.task_spec is None or not ctx.task_spec

    @pytest.mark.asyncio
    async def test_empty_response_returns_failure(self):
        """Empty LLM response → AgentResult(success=False)."""
        agent = _make_agent("")
        ctx = AgentContext(run_id="r2", user_intent="classify emails")
        result = await agent.run(ctx)

        assert result.success is False

    @pytest.mark.asyncio
    async def test_partial_json_missing_task_type_returns_failure(self):
        """JSON missing required 'task_type' field → parse/validation fails."""
        bad_json = json.dumps({
            "num_labels": 2,
            "confidence": 0.9,
        })
        agent = _make_agent(bad_json)
        ctx = AgentContext(run_id="r3", user_intent="classify emails")
        result = await agent.run(ctx)

        assert result.success is False

    @pytest.mark.asyncio
    async def test_markdown_wrapped_json_is_tolerated_or_fails_gracefully(self):
        """
        LLM sometimes wraps JSON in ```json ... ``` fences.
        The agent should either strip them (works) or fail gracefully (success=False).
        Either is acceptable — it must NOT raise an unhandled exception.
        """
        wrapped = "```json\n" + _valid_json() + "\n```"
        agent = _make_agent(wrapped)
        ctx = AgentContext(run_id="r4", user_intent="classify support tickets")
        result = await agent.run(ctx)  # must not raise

        # Either success (agent strips fences) or graceful failure — not an exception
        assert isinstance(result.success, bool)


# ── Successful parse ──────────────────────────────────────────────────────────

class TestIntentAgentSuccessfulParse:
    @pytest.mark.asyncio
    async def test_context_task_spec_populated(self):
        """After a successful run, context.task_spec must be populated."""
        agent = _make_agent(_valid_json())
        ctx = AgentContext(run_id="r5", user_intent="classify support tickets by urgency")
        result = await agent.run(ctx)

        assert result.success is True
        assert ctx.task_spec is not None
        assert ctx.task_spec.get("task_type") == "text_classification"

    @pytest.mark.asyncio
    async def test_result_output_contains_task_type(self):
        """result.output must carry back the full parsed spec dict."""
        agent = _make_agent(_valid_json(task_type="text_classification"))
        ctx = AgentContext(run_id="r6", user_intent="classify sentiment")
        result = await agent.run(ctx)

        assert result.output.get("task_type") == "text_classification"

    @pytest.mark.asyncio
    async def test_result_output_contains_base_model_hint(self):
        """base_model_hint must appear in result.output."""
        agent = _make_agent(_valid_json(base_model="roberta-base"))
        ctx = AgentContext(run_id="r7", user_intent="classify product reviews")
        result = await agent.run(ctx)

        assert result.output.get("base_model_hint") == "roberta-base"

    @pytest.mark.asyncio
    async def test_success_message_mentions_task_type(self):
        """On success, the message should describe the task type."""
        agent = _make_agent(_valid_json(task_type="text_classification"))
        ctx = AgentContext(run_id="r8", user_intent="classify emails")
        result = await agent.run(ctx)

        assert result.success is True
        # Message should mention the task label
        assert "classification" in result.message.lower() or "text" in result.message.lower()

    @pytest.mark.asyncio
    async def test_next_agent_is_data_on_high_confidence(self):
        """High confidence → next_agent == 'Data'."""
        agent = _make_agent(_valid_json(confidence=0.9))
        ctx = AgentContext(run_id="r9", user_intent="classify support tickets")
        result = await agent.run(ctx)

        assert result.next_agent == "Data"


# ── NER task type ─────────────────────────────────────────────────────────────

class TestIntentAgentNER:
    @pytest.mark.asyncio
    async def test_ner_task_type_parsed(self):
        """NER intent → task_type='token_classification' stored in context."""
        agent = _make_agent(_valid_json(
            task_type="token_classification",
            label_names=["PER", "ORG", "LOC"],
            input_col="tokens",
            label_col="ner_tags",
            base_model="dslim/bert-base-NER",
            confidence=0.9,
        ))
        ctx = AgentContext(run_id="r10", user_intent="extract named entities from news articles")
        result = await agent.run(ctx)

        assert result.success is True
        assert ctx.task_spec["task_type"] == "token_classification"
        assert ctx.task_spec["label_names"] == ["PER", "ORG", "LOC"]
        assert ctx.task_spec["input_column"] == "tokens"
        assert ctx.task_spec["label_column"] == "ner_tags"

    @pytest.mark.asyncio
    async def test_ner_routes_to_data_agent(self):
        """Token classification with high confidence should also route to Data."""
        agent = _make_agent(_valid_json(
            task_type="token_classification",
            confidence=0.88,
        ))
        ctx = AgentContext(run_id="r11", user_intent="NER on clinical notes")
        result = await agent.run(ctx)

        assert result.next_agent == "Data"


# ── System prompt cache eligibility ──────────────────────────────────────────

class TestIntentAgentSystemPrompt:
    def test_system_prompt_meets_cache_minimum(self):
        """
        Anthropic's prompt caching requires ≥ 1 024 tokens in the system prompt.
        We use a character-count proxy: 1 token ≈ 4 characters (conservative
        lower-bound for English text).  1 024 tokens ≈ 4 096 characters.
        """
        from agents.intent import SYSTEM

        # 4 096 chars ≈ 1 024 tokens at 4 chars/token (conservative lower bound).
        min_chars = 4_096
        assert len(SYSTEM) >= min_chars, (
            f"SYSTEM prompt is only {len(SYSTEM)} characters (~{len(SYSTEM)//4} tokens) — "
            f"under the 1 024-token cache minimum. Expand it so cache_system=True "
            "actually reduces API costs."
        )