cleveragents-core/features/steps/semantic_context_search_steps.py

"""Step definitions for semantic context search feature tests."""

from __future__ import annotations

from behave import given, then, when

from cleveragents.application.services.embedding_provider import (
    MockEmbeddingProvider,
    SimpleWordEmbeddingProvider,
    cosine_similarity,
)
from cleveragents.domain.models.core.context_fragment import (
    ContextBudget,
    ContextFragment,
    FragmentProvenance,
)


@given("I have a semantic context strategy with embeddings")
def step_have_semantic_strategy(context):
    """Initialize a semantic context strategy."""
    context.embedding_provider = MockEmbeddingProvider(embedding_dim=10)
    context.fragments = []


@given("I have a mock embedding provider")
def step_have_mock_provider(context):
    """Initialize a mock embedding provider."""
    context.embedding_provider = MockEmbeddingProvider(embedding_dim=10)


@when("I embed the text {text}")
def step_embed_text(context, text):
    """Embed a text fragment."""
    context.embedding = context.embedding_provider.embed(text)


@then("the embedding should have {dim:d} dimensions")
def step_check_embedding_dimension(context, dim):
    """Verify embedding dimension."""
    assert len(context.embedding) == dim, (
        f"Expected {dim} dimensions, got {len(context.embedding)}"
    )


@then("the embedding should be a valid vector")
def step_check_valid_vector(context):
    """Verify embedding is a valid vector."""
    assert isinstance(context.embedding, (list, tuple))
    assert all(isinstance(x, (int, float)) for x in context.embedding)
    assert len(context.embedding) > 0


@given("I have two text embeddings")
def step_have_two_embeddings(context):
    """Create two text embeddings."""
    context.embedding1 = context.embedding_provider.embed("Python function")
    context.embedding2 = context.embedding_provider.embed("Python class")


@when("I compute their cosine similarity")
def step_compute_similarity(context):
    """Compute cosine similarity."""
    context.similarity = cosine_similarity(context.embedding1, context.embedding2)


@then("the similarity should be between -1 and 1")
def step_check_similarity_range(context):
    """Verify similarity is in valid range."""
    assert -1 <= context.similarity <= 1, (
        f"Similarity {context.similarity} out of range"
    )


@given("I have context fragments with content:")
def step_have_fragments_with_content(context):
    """Create context fragments from table."""
    context.fragments = []
    for row in context.table:
        content = row["content"]
        embedding = context.embedding_provider.embed(content)
        context.fragments.append(
            {
                "content": content,
                "embedding": embedding,
            }
        )


@given("I have a query {query}")
def step_have_query(context, query):
    """Set the query."""
    context.query = query
    context.query_embedding = context.embedding_provider.embed(query)


@when("I rank fragments by semantic similarity to the query")
def step_rank_fragments(context):
    """Rank fragments by similarity."""
    similarities = []
    for frag in context.fragments:
        sim = cosine_similarity(context.query_embedding, frag["embedding"])
        similarities.append((frag, sim))

    similarities.sort(key=lambda x: x[1], reverse=True)
    context.ranked_fragments = similarities


@then("the Python fragments should rank higher than JavaScript")
def step_check_python_ranking(context):
    """Verify Python fragments rank higher."""
    python_sims = [
        sim for frag, sim in context.ranked_fragments if "Python" in frag["content"]
    ]
    js_sims = [
        sim for frag, sim in context.ranked_fragments if "JavaScript" in frag["content"]
    ]

    if python_sims and js_sims:
        assert min(python_sims) >= max(js_sims), "Python fragments should rank higher"


@given("I have a minimum similarity threshold of {threshold:f}")
def step_have_threshold(context, threshold):
    """Set similarity threshold."""
    context.threshold = threshold


@when("I filter fragments by similarity threshold")
def step_filter_by_threshold(context):
    """Filter fragments by threshold.

    Computes similarity inline so the step works whether or not a prior
    "rank fragments" step ran. Scenarios that filter directly (without an
    intermediate ranking step) would otherwise hit an AttributeError on
    ``context.ranked_fragments``.
    """
    scored = [
        (frag, cosine_similarity(context.query_embedding, frag["embedding"]))
        for frag in context.fragments
    ]
    context.filtered_fragments = [
        (frag, sim) for frag, sim in scored if sim >= context.threshold
    ]


@then("only semantically similar fragments should be included")
def step_check_filtered_fragments(context):
    """Verify filtered fragments meet threshold."""
    for _frag, sim in context.filtered_fragments:
        assert sim >= context.threshold, (
            f"Fragment similarity {sim} below threshold {context.threshold}"
        )


@given("I have a semantic context strategy")
def step_have_strategy(context):
    """Initialize semantic strategy."""
    context.embedding_provider = SimpleWordEmbeddingProvider(vocab_size=100)


@given("I have context fragments:")
def step_have_context_fragments(context):
    """Create context fragments from table."""
    context.fragments = []
    for row in context.table:
        frag = ContextFragment(
            uko_node=row["uko_node"],
            content=row["content"],
            token_count=int(row["token_count"]),
            relevance_score=0.5,
            detail_depth=1,
            tier="hot",
            provenance=FragmentProvenance(resource_uri=row["uko_node"]),
        )
        context.fragments.append(frag)


@given("I have a context budget of {tokens:d} tokens")
def step_have_budget(context, tokens):
    """Set context budget."""
    context.budget = ContextBudget(max_tokens=tokens, reserved_tokens=0)


@when("I assemble semantic context with query {query}")
def step_assemble_context(context, query):
    """Assemble context with query."""
    context.query = query
    query_embedding = context.embedding_provider.embed(query)

    # Score fragments by similarity
    scored = []
    for frag in context.fragments:
        frag_embedding = context.embedding_provider.embed(frag.content)
        sim = cosine_similarity(query_embedding, frag_embedding)
        scored.append((frag, sim))

    # Sort by similarity
    scored.sort(key=lambda x: x[1], reverse=True)

    # Pack within budget
    context.selected_fragments = []
    total_tokens = 0
    for frag, _sim in scored:
        if total_tokens + frag.token_count <= context.budget.max_tokens:
            context.selected_fragments.append(frag)
            total_tokens += frag.token_count


@then("the selected fragments should include Python files")
def step_check_python_files(context):
    """Verify Python files are selected."""
    python_files = [f for f in context.selected_fragments if ".py" in f.uko_node]
    assert len(python_files) > 0, "No Python files selected"


@then("the selected fragments should be ranked by relevance")
def step_check_ranking(context):
    """Verify fragments are ranked by relevance."""
    assert len(context.selected_fragments) > 0, "No fragments selected"


@given("I have an embedding provider configuration")
def step_have_config(context):
    """Initialize embedding provider configuration.

    Stored under ``embedding_config`` rather than ``config`` because Behave's
    ``Context`` reserves ``config`` for its own Configuration object — setting
    ``context.config`` raises ``KeyError`` from Behave's scope-tracking
    ``__setattr__``.
    """
    context.embedding_config = {
        "provider_type": "simple_word",
        "vocab_size": 100,
    }


@when("I create a semantic strategy with the configuration")
def step_create_strategy_with_config(context):
    """Create strategy with configuration."""
    if context.embedding_config["provider_type"] == "simple_word":
        context.strategy_provider = SimpleWordEmbeddingProvider(
            vocab_size=context.embedding_config["vocab_size"]
        )


@then("the strategy should use the configured embedding provider")
def step_check_configured_provider(context):
    """Verify strategy uses configured provider."""
    assert context.strategy_provider is not None
    assert context.strategy_provider.embedding_dimension > 0


@then("the strategy should report semantic search capability")
def step_check_capability(context):
    """Verify strategy reports semantic capability."""
    assert hasattr(context.strategy_provider, "embedding_dimension")
    assert context.strategy_provider.embedding_dimension > 0