op_agent/test_extract_code_tool.py

"""
Test script for extract_first_code as an LLM-callable tool.

Demonstrates:
  1. Direct handler call (old behaviour, still works)
  2. Tool-schema validation against Anthropic / OpenAI shapes
  3. Simulated tool_use round-trip — model output → tool call → result
  4. Full Anthropic API round-trip with a real LLM (optional, skipped if no key)

Usage:
    python test_extract_code_tool.py           # unit tests only
    python test_extract_code_tool.py --live    # also test against a real LLM
"""

from __future__ import annotations

import json
import os
import sys

import pytest

# Make agent/ importable from repo root
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from agent.tool_extract_code import (
    EXTRACT_CODE_TOOL_ANTHROPIC,
    EXTRACT_CODE_TOOL_OPENAI,
    extract_first_code,
    handle_tool_call,
)

# ── Fixtures ─────────────────────────────────────────────────────────────────


SAMPLE_WITH_PYTHON = """\
Here is a Triton kernel:

```python
import triton
import triton.language as tl

@triton.jit
def softmax_kernel(x_ptr, y_ptr, N):
    ...
```

This kernel computes softmax.
"""

SAMPLE_WITH_CPP = """\
```cpp
__global__ void vec_add(float* a, float* b, float* c, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) c[i] = a[i] + b[i];
}
```
"""

SAMPLE_NO_FENCE = "just a plain string without any code block"

SAMPLE_MULTIPLE = """\
```python
print("first")
```
Some text in between.
```python
print("second")
```
"""


# ── Direct handler tests ─────────────────────────────────────────────────────


def test_extract_python():
    result = extract_first_code(SAMPLE_WITH_PYTHON, ["python", "cpp"])
    assert result.startswith("import triton")
    assert "```" not in result


def test_extract_cpp():
    result = extract_first_code(SAMPLE_WITH_CPP, ["python", "cpp"])
    assert "vec_add" in result
    assert result.startswith("__global__")


def test_no_fence_returns_original():
    result = extract_first_code(SAMPLE_NO_FENCE, ["python"])
    assert result == SAMPLE_NO_FENCE


def test_multiple_fences_extracts_first():
    result = extract_first_code(SAMPLE_MULTIPLE, ["python"])
    assert 'print("first")' in result
    assert 'print("second")' not in result


def test_language_tag_stripped():
    result = extract_first_code("```python\nx = 1\n```", ["python"])
    assert result == "x = 1"


def test_unknown_language_tag_kept():
    result = extract_first_code("```rust\nfn main() {}\n```", ["python", "cpp"])
    assert result == "rust\nfn main() {}"


# ── Tool dispatch tests ──────────────────────────────────────────────────────


def test_handle_tool_call():
    result = handle_tool_call(
        "extract_first_code",
        {"output_string": SAMPLE_WITH_PYTHON, "code_language_types": ["python"]},
    )
    assert result.startswith("import triton")


def test_handle_unknown_tool_raises():
    with pytest.raises(ValueError, match="Unknown tool"):
        handle_tool_call("nonexistent", {})


# ── Schema validation tests ──────────────────────────────────────────────────


def test_anthropic_schema_is_valid():
    """Anthropic tool schema must have name, description, input_schema."""
    assert EXTRACT_CODE_TOOL_ANTHROPIC["name"] == "extract_first_code"
    assert "description" in EXTRACT_CODE_TOOL_ANTHROPIC
    schema = EXTRACT_CODE_TOOL_ANTHROPIC["input_schema"]
    assert schema["type"] == "object"
    assert "output_string" in schema["properties"]
    assert "code_language_types" in schema["properties"]
    assert set(schema["required"]) == {"output_string", "code_language_types"}


def test_openai_schema_is_valid():
    """OpenAI tool schema must have type: function and a function sub-object."""
    assert EXTRACT_CODE_TOOL_OPENAI["type"] == "function"
    fn = EXTRACT_CODE_TOOL_OPENAI["function"]
    assert fn["name"] == "extract_first_code"
    assert "parameters" in fn


# ── Simulated round-trip ─────────────────────────────────────────────────────


def simulate_anthropic_tool_round_trip(model_output: str) -> str:
    """
    Mimic what happens when the LLM returns a tool_use block:
      1. LLM emits text + tool_use(content=model_output)
      2. Client code dispatches to handle_tool_call
      3. Client code sends back a tool_result
    """
    return handle_tool_call(
        "extract_first_code",
        {"output_string": model_output, "code_language_types": ["python", "cpp"]},
    )


def test_simulated_round_trip():
    kernel = simulate_anthropic_tool_round_trip(SAMPLE_WITH_PYTHON)
    assert "import triton" in kernel
    assert "```" not in kernel


# ── Live LLM test (optional) ─────────────────────────────────────────────────

def test_live_anthropic_tool_call():
    """Real round-trip against the Anthropic API. Skipped without --live."""
    if "--live" not in sys.argv:
        pytest.skip("pass --live to run against a real LLM")

    import anthropic  # type: ignore[import-not-found]

    api_key = os.environ.get("ANTHROPIC_API_KEY") or os.environ.get("ANTHROPIC_AUTH_TOKEN")
    if not api_key:
        pytest.skip("No ANTHROPIC_API_KEY set")

    client = anthropic.Anthropic(api_key=api_key)

    response = client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=2048,
        tools=[EXTRACT_CODE_TOOL_ANTHROPIC],
        messages=[
            {
                "role": "user",
                "content": (
                    "Write a short python function that adds two numbers, "
                    "then call extract_first_code on your own response "
                    "to isolate the code."
                ),
            }
        ],
    )

    # Collect text and any tool_use
    tool_uses = [b for b in response.content if b.type == "tool_use"]
    assert len(tool_uses) > 0, "LLM should have called extract_first_code"

    for block in tool_uses:
        result = handle_tool_call(block.name, block.input)
        assert "def " in result, f"Unexpected extracted code:\n{result}"
        print(f"\n[OK] LLM called tool, extracted:\n{result}")


if __name__ == "__main__":
    sys.exit(pytest.main([__file__, "-v"] + sys.argv[1:]))