cleveragents-core/features/steps/crp_models_steps.py

"""Step definitions for the CRP Domain Models feature."""

from __future__ import annotations

from typing import Any

from behave import given, then, when
from pydantic import ValidationError

from cleveragents.domain.models.acms.crp import (
    AssembledContext,
    ContextBudget,
    ContextFragment,
    ContextRequest,
    DetailLevelMap,
    FragmentProvenance,
)
from cleveragents.domain.models.core.project import TemporalScope
from cleveragents.skills.builtins.context_ops import (
    _handle_get_context_budget,
    _handle_query_history,
    _handle_request_context,
    register_skill_context_tools,
)
from cleveragents.tool.registry import ToolRegistry

__all__: list[str] = []


# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------


def _make_provenance(
    resource_uri: str = "uko-py:module/test",
) -> FragmentProvenance:
    """Create a simple provenance for test fragments."""
    return FragmentProvenance(resource_uri=resource_uri)


# ---------------------------------------------------------------------------
# DetailLevelMap
# ---------------------------------------------------------------------------


@given('a DetailLevelMap for domain "{domain}" with max depth {max_depth:d}')
def step_given_detail_level_map(context: Any, domain: str, max_depth: int) -> None:
    context.detail_level_map = DetailLevelMap(domain=domain, max_depth=max_depth)


@given('the map has level "{name}" at depth {depth:d}')
def step_given_map_has_level(context: Any, name: str, depth: int) -> None:
    context.detail_level_map.register(name, depth)


@then('the map should resolve "{name}" to {expected:d}')
def step_then_map_resolve_name(context: Any, name: str, expected: int) -> None:
    actual: int = context.detail_level_map.resolve(name)
    assert actual == expected, f"Expected {expected}, got {actual}"


@then("the map should resolve integer {value:d} to {expected:d}")
def step_then_map_resolve_int(context: Any, value: int, expected: int) -> None:
    actual: int = context.detail_level_map.resolve(value)
    assert actual == expected, f"Expected {expected}, got {actual}"


@given('a parent DetailLevelMap for domain "{domain}" with max depth {max_depth:d}')
def step_given_parent_map(context: Any, domain: str, max_depth: int) -> None:
    context.parent_detail_level_map = DetailLevelMap(domain=domain, max_depth=max_depth)


@given('the parent map has level "{name}" at depth {depth:d}')
def step_given_parent_level(context: Any, name: str, depth: int) -> None:
    context.parent_detail_level_map.register(name, depth)


@given('a child DetailLevelMap for domain "{domain}" with max depth {max_depth:d}')
def step_given_child_map(context: Any, domain: str, max_depth: int) -> None:
    context.child_detail_level_map = DetailLevelMap(
        domain=domain,
        max_depth=max_depth,
        parent=context.parent_detail_level_map,
    )


@given('the child map has level "{name}" at depth {depth:d}')
def step_given_child_level(context: Any, name: str, depth: int) -> None:
    context.child_detail_level_map.register(name, depth)


@then('the child map should resolve "{name}" to {expected:d}')
def step_then_child_resolve(context: Any, name: str, expected: int) -> None:
    actual: int = context.child_detail_level_map.resolve(name)
    assert actual == expected, f"Expected {expected}, got {actual}"


@then('resolving "{name}" should raise ValueError')
def step_then_resolve_raises(context: Any, name: str) -> None:
    try:
        context.detail_level_map.resolve(name)
        raise AssertionError(f"Expected ValueError for '{name}'")
    except ValueError:
        pass


@when('I register level "{name}" at depth {depth:d}')
def step_when_register_level(context: Any, name: str, depth: int) -> None:
    context.detail_level_map.register(name, depth)


@then(
    'creating a DetailLevelMap with level "{name}" at depth {depth:d} '
    "should raise ValueError"
)
def step_then_bad_level_raises(context: Any, name: str, depth: int) -> None:
    try:
        DetailLevelMap(domain="test:", max_depth=9, levels={name: depth})
        raise AssertionError("Expected ValueError")
    except (ValueError, ValidationError):
        pass


@then('registering level "{name}" at depth {depth:d} should raise ValueError')
def step_then_register_bad_raises(context: Any, name: str, depth: int) -> None:
    try:
        context.detail_level_map.register(name, depth)
        raise AssertionError("Expected ValueError")
    except ValueError:
        pass


@given(
    'a DetailLevelMap for domain "{domain}" with max depth {max_depth:d} '
    'and levels "{name}" at {depth:d}'
)
def step_given_detail_level_map_with_levels(
    context: Any, domain: str, max_depth: int, name: str, depth: int
) -> None:
    context.detail_level_map = DetailLevelMap(
        domain=domain, max_depth=max_depth, levels={name: depth}
    )


@then('the child effective_levels should contain "{name}" at {expected:d}')
def step_then_effective_levels(context: Any, name: str, expected: int) -> None:
    levels: dict[str, int] = context.child_detail_level_map.effective_levels()
    assert name in levels, f"'{name}' not in effective_levels"
    assert levels[name] == expected, (
        f"Expected {expected} for '{name}', got {levels[name]}"
    )


# ---------------------------------------------------------------------------
# FragmentProvenance
# ---------------------------------------------------------------------------


@given('a FragmentProvenance with resource URI "{uri}"')
def step_given_provenance(context: Any, uri: str) -> None:
    context.provenance = FragmentProvenance(resource_uri=uri)


@then('the provenance resource_uri should be "{expected}"')
def step_then_provenance_uri(context: Any, expected: str) -> None:
    assert context.provenance.resource_uri == expected


@then("the provenance location should be empty")
def step_then_provenance_location_empty(context: Any) -> None:
    assert context.provenance.location == ""


@then("the provenance strategy should be empty")
def step_then_provenance_strategy_empty(context: Any) -> None:
    assert context.provenance.strategy == ""


@then("creating a FragmentProvenance with empty resource_uri should raise ValueError")
def step_then_provenance_empty_uri(context: Any) -> None:
    try:
        FragmentProvenance(resource_uri="")
        raise AssertionError("Expected ValueError")
    except (ValueError, ValidationError):
        pass


# ---------------------------------------------------------------------------
# ContextBudget
# ---------------------------------------------------------------------------


@given("a ContextBudget with max_tokens {max_tok:d} and reserved_tokens {res:d}")
def step_given_budget(context: Any, max_tok: int, res: int) -> None:
    context.budget = ContextBudget(max_tokens=max_tok, reserved_tokens=res)


@then("the budget available_tokens should be {expected:d}")
def step_then_budget_available(context: Any, expected: int) -> None:
    assert context.budget.available_tokens == expected


@then("the budget max_tokens should be {expected:d}")
def step_then_budget_max(context: Any, expected: int) -> None:
    assert context.budget.max_tokens == expected


@then("the budget reserved_tokens should be {expected:d}")
def step_then_budget_reserved(context: Any, expected: int) -> None:
    assert context.budget.reserved_tokens == expected


@then("creating a ContextBudget with max_tokens {max_tok:d} should raise ValueError")
def step_then_budget_bad_max(context: Any, max_tok: int) -> None:
    try:
        ContextBudget(max_tokens=max_tok)
        raise AssertionError("Expected ValueError")
    except (ValueError, ValidationError):
        pass


@then(
    "creating a ContextBudget with max_tokens {max_tok:d} "
    "and reserved_tokens {res:d} should raise ValueError"
)
def step_then_budget_bad_reserved(context: Any, max_tok: int, res: int) -> None:
    try:
        ContextBudget(max_tokens=max_tok, reserved_tokens=res)
        raise AssertionError("Expected ValueError")
    except (ValueError, ValidationError):
        pass


# ---------------------------------------------------------------------------
# ContextFragment
# ---------------------------------------------------------------------------


@given('a ContextFragment with uko_node "{node}" and content "{content}"')
def step_given_fragment(context: Any, node: str, content: str) -> None:
    context.fragment_node = node
    context.fragment_content = content
    context.fragment_detail_depth = 0
    context.fragment_token_count = 0
    context.fragment_relevance = 0.5


@given("the fragment has detail_depth {depth:d} and token_count {tokens:d}")
def step_given_fragment_depth_tokens(context: Any, depth: int, tokens: int) -> None:
    context.fragment_detail_depth = depth
    context.fragment_token_count = tokens


@given("the fragment has relevance_score {score:g}")
def step_given_fragment_relevance(context: Any, score: float) -> None:
    context.fragment_relevance = score
    context.fragment = ContextFragment(
        uko_node=context.fragment_node,
        content=context.fragment_content,
        detail_depth=context.fragment_detail_depth,
        token_count=context.fragment_token_count,
        relevance_score=context.fragment_relevance,
        provenance=_make_provenance(),
    )


@then('the fragment uko_node should be "{expected}"')
def step_then_fragment_node(context: Any, expected: str) -> None:
    assert context.fragment.uko_node == expected


@then("the fragment token_count should be {expected:d}")
def step_then_fragment_tokens(context: Any, expected: int) -> None:
    assert context.fragment.token_count == expected


@then("the fragment relevance_score should be {expected:g}")
def step_then_fragment_relevance(context: Any, expected: float) -> None:
    assert context.fragment.relevance_score == expected


@then("the fragment detail_depth should be {expected:d}")
def step_then_fragment_depth(context: Any, expected: int) -> None:
    assert context.fragment.detail_depth == expected


@then("creating a ContextFragment with token_count {count:d} should raise ValueError")
def step_then_fragment_bad_tokens(context: Any, count: int) -> None:
    try:
        ContextFragment(
            uko_node="uko-py:test",
            content="test",
            detail_depth=0,
            token_count=count,
            relevance_score=0.5,
            provenance=_make_provenance(),
        )
        raise AssertionError("Expected ValueError")
    except (ValueError, ValidationError):
        pass


@then(
    "creating a ContextFragment with relevance_score {score:g} should raise ValueError"
)
def step_then_fragment_bad_relevance(context: Any, score: float) -> None:
    try:
        ContextFragment(
            uko_node="uko-py:test",
            content="test",
            detail_depth=0,
            token_count=10,
            relevance_score=score,
            provenance=_make_provenance(),
        )
        raise AssertionError("Expected ValueError")
    except (ValueError, ValidationError):
        pass


@then("creating a ContextFragment with detail_depth {depth:d} should raise ValueError")
def step_then_fragment_bad_depth(context: Any, depth: int) -> None:
    try:
        ContextFragment(
            uko_node="uko-py:test",
            content="test",
            detail_depth=depth,
            token_count=10,
            relevance_score=0.5,
            provenance=_make_provenance(),
        )
        raise AssertionError("Expected ValueError")
    except (ValueError, ValidationError):
        pass


@then("creating a ContextFragment with empty uko_node should raise ValueError")
def step_then_fragment_empty_node(context: Any) -> None:
    try:
        ContextFragment(
            uko_node="",
            content="test",
            detail_depth=0,
            token_count=10,
            relevance_score=0.5,
            provenance=_make_provenance(),
        )
        raise AssertionError("Expected ValueError")
    except (ValueError, ValidationError):
        pass


# ---------------------------------------------------------------------------
# ContextRequest
# ---------------------------------------------------------------------------


@given('a ContextRequest with purpose "{purpose}"')
def step_given_request(context: Any, purpose: str) -> None:
    context.ctx_request = ContextRequest(purpose=purpose)


@given('a full ContextRequest with query "{query}" and focus "{focus}"')
def step_given_full_request(context: Any, query: str, focus: str) -> None:
    context.ctx_request = ContextRequest(
        query=query,
        focus=[focus],
    )


@given('the request has breadth {breadth:d} and depth "{depth}"')
def step_given_request_breadth_depth_str(
    context: Any, breadth: int, depth: str
) -> None:
    context.ctx_request = context.ctx_request.model_copy(
        update={"breadth": breadth, "depth": depth}
    )


@given('the request has priority {priority:g} and purpose "{purpose}"')
def step_given_request_priority_purpose(
    context: Any, priority: float, purpose: str
) -> None:
    context.ctx_request = context.ctx_request.model_copy(
        update={"priority": priority, "purpose": purpose}
    )


@given('a ContextRequest with depth "{depth}"')
def step_given_request_depth_str(context: Any, depth: str) -> None:
    context.ctx_request = ContextRequest(depth=depth)


@given("a ContextRequest with breadth {breadth:d}")
def step_given_request_breadth(context: Any, breadth: int) -> None:
    context.ctx_request = ContextRequest(breadth=breadth)


@given("a ContextRequest with max_tokens {max_tok:d}")
def step_given_request_max_tokens(context: Any, max_tok: int) -> None:
    context.ctx_request = ContextRequest(max_tokens=max_tok)


@then('the request purpose should be "{expected}"')
def step_then_request_purpose(context: Any, expected: str) -> None:
    assert context.ctx_request.purpose == expected


@then("the request breadth should be {expected:d}")
def step_then_request_breadth(context: Any, expected: int) -> None:
    assert context.ctx_request.breadth == expected


@then("the request depth should be {expected:d}")
def step_then_request_depth_int(context: Any, expected: int) -> None:
    assert context.ctx_request.depth == expected


@then('the request depth should be "{expected}"')
def step_then_request_depth_str(context: Any, expected: str) -> None:
    assert context.ctx_request.depth == expected


@then("the request depth_gradient should be true")
def step_then_request_gradient_true(context: Any) -> None:
    assert context.ctx_request.depth_gradient is True


@then("the request priority should be {expected:g}")
def step_then_request_priority(context: Any, expected: float) -> None:
    assert context.ctx_request.priority == expected


@then('the request temporal should be "{expected}"')
def step_then_request_temporal(context: Any, expected: str) -> None:
    assert context.ctx_request.temporal == TemporalScope(expected)


@then("the request query should be None")
def step_then_request_query_none(context: Any) -> None:
    assert context.ctx_request.query is None


@then('the request query should be "{expected}"')
def step_then_request_query(context: Any, expected: str) -> None:
    assert context.ctx_request.query == expected


@then("the request max_tokens should be None")
def step_then_request_max_tokens_none(context: Any) -> None:
    assert context.ctx_request.max_tokens is None


@then("the request max_tokens should be {expected:d}")
def step_then_request_max_tokens(context: Any, expected: int) -> None:
    assert context.ctx_request.max_tokens == expected


@then("creating a ContextRequest with breadth {breadth:d} should raise ValueError")
def step_then_request_bad_breadth(context: Any, breadth: int) -> None:
    try:
        ContextRequest(breadth=breadth)
        raise AssertionError("Expected ValueError")
    except (ValueError, ValidationError):
        pass


@then("creating a ContextRequest with depth {depth:d} should raise ValueError")
def step_then_request_bad_depth(context: Any, depth: int) -> None:
    try:
        ContextRequest(depth=depth)
        raise AssertionError("Expected ValueError")
    except (ValueError, ValidationError):
        pass


@then("creating a ContextRequest with priority {priority:g} should raise ValueError")
def step_then_request_bad_priority(context: Any, priority: float) -> None:
    try:
        ContextRequest(priority=priority)
        raise AssertionError("Expected ValueError")
    except (ValueError, ValidationError):
        pass


@then("creating a ContextRequest with max_tokens {max_tok:d} should raise ValueError")
def step_then_request_bad_max_tokens(context: Any, max_tok: int) -> None:
    try:
        ContextRequest(max_tokens=max_tok)
        raise AssertionError("Expected ValueError")
    except (ValueError, ValidationError):
        pass


# ---------------------------------------------------------------------------
# AssembledContext
# ---------------------------------------------------------------------------


@given("an AssembledContext with total_tokens {total:d} and budget_used {used:g}")
def step_given_assembled(context: Any, total: int, used: float) -> None:
    context.assembled_total = total
    context.assembled_used = used


@given('the assembled context has hash "{hash_val}"')
def step_given_assembled_hash(context: Any, hash_val: str) -> None:
    context.assembled = AssembledContext(
        total_tokens=context.assembled_total,
        budget_used=context.assembled_used,
        context_hash=hash_val,
    )


@then("the assembled total_tokens should be {expected:d}")
def step_then_assembled_total(context: Any, expected: int) -> None:
    assert context.assembled.total_tokens == expected


@then("the assembled budget_used should be {expected:g}")
def step_then_assembled_used(context: Any, expected: float) -> None:
    assert context.assembled.budget_used == expected


@then('the assembled context_hash should be "{expected}"')
def step_then_assembled_hash(context: Any, expected: str) -> None:
    assert context.assembled.context_hash == expected


@then("creating an AssembledContext with budget_used {used:g} should raise ValueError")
def step_then_assembled_bad_used(context: Any, used: float) -> None:
    try:
        AssembledContext(
            total_tokens=100,
            budget_used=used,
            context_hash="test",
        )
        raise AssertionError("Expected ValueError")
    except (ValueError, ValidationError):
        pass


@then(
    "creating an AssembledContext with total_tokens {total:d} should raise ValueError"
)
def step_then_assembled_bad_total(context: Any, total: int) -> None:
    try:
        AssembledContext(
            total_tokens=total,
            budget_used=0.5,
            context_hash="test",
        )
        raise AssertionError("Expected ValueError")
    except (ValueError, ValidationError):
        pass


# ---------------------------------------------------------------------------
# Context skill tool stubs
# ---------------------------------------------------------------------------


@when("I call the request_context tool handler")
def step_when_call_request_context(context: Any) -> None:
    context.tool_error = None
    try:
        _handle_request_context({"purpose": "test"})
    except NotImplementedError as exc:
        context.tool_error = exc


@when("I call the query_history tool handler")
def step_when_call_query_history(context: Any) -> None:
    context.tool_error = None
    try:
        _handle_query_history({"query": "test"})
    except NotImplementedError as exc:
        context.tool_error = exc


@when("I call the get_context_budget tool handler")
def step_when_call_get_budget(context: Any) -> None:
    context.tool_error = None
    try:
        _handle_get_context_budget({})
    except NotImplementedError as exc:
        context.tool_error = exc


@then('it should raise NotImplementedError with message containing "{text}"')
def step_then_not_impl_error(context: Any, text: str) -> None:
    assert context.tool_error is not None, "Expected NotImplementedError"
    assert text in str(context.tool_error), (
        f"Expected '{text}' in '{context.tool_error}'"
    )


@given("a CRP tool registry with context tools registered")
def step_given_crp_registry_with_context(context: Any) -> None:
    context.crp_tool_registry = ToolRegistry()
    register_skill_context_tools(context.crp_tool_registry)


@then('the CRP registry should contain tool "{name}"')
def step_then_crp_registry_has_tool(context: Any, name: str) -> None:
    spec = context.crp_tool_registry.get(name)
    assert spec is not None, f"Tool '{name}' not found in CRP registry"


@then('CRP tool "{name}" should have read_only capability')
def step_then_crp_tool_read_only(context: Any, name: str) -> None:
    spec = context.crp_tool_registry.get(name)
    assert spec is not None, f"Tool '{name}' not found"
    assert spec.capabilities.read_only is True, f"Tool '{name}' is not read-only"


@then('CRP tool "{name}" should require "{field}" in input schema')
def step_then_crp_tool_requires_field(context: Any, name: str, field: str) -> None:
    spec = context.crp_tool_registry.get(name)
    assert spec is not None, f"Tool '{name}' not found"
    required: list[str] = spec.input_schema.get("required", [])
    assert field in required, f"Field '{field}' not in required fields: {required}"