cleveragents-core/robot/helper_context_strategies.py

"""Robot Framework integration helper for context strategies batch 1.

Each ``_cmd_*`` function exercises one strategy end-to-end and prints a
sentinel string on success.  The Robot ``.robot`` file invokes this script
as a subprocess via ``Run Process`` and asserts exit-code + sentinel.
"""

from __future__ import annotations

import sys
from pathlib import Path

# Bootstrap src/ so domain imports resolve when executed standalone.
_src = str(Path(__file__).resolve().parent.parent / "src")
if _src not in sys.path:
    sys.path.insert(0, _src)

from cleveragents.application.services.acms_service import ACMSPipeline  # noqa: E402
from cleveragents.application.services.context_strategies import (  # noqa: E402
    BreadthDepthNavigatorStrategy,
    SemanticEmbeddingStrategy,
    SimpleKeywordStrategy,
)
from cleveragents.domain.models.core.context_fragment import (  # noqa: E402
    ContextBudget,
    ContextFragment,
    FragmentProvenance,
)


def _make_frag(
    uko_node: str,
    content: str,
    *,
    score: float = 0.5,
    tokens: int = 20,
    depth: int = 3,
) -> ContextFragment:
    return ContextFragment(
        uko_node=uko_node,
        content=content,
        relevance_score=score,
        token_count=tokens,
        detail_depth=depth,
        provenance=FragmentProvenance(resource_uri=uko_node),
    )


# ---------------------------------------------------------------------------
# Command functions
# ---------------------------------------------------------------------------


def _cmd_simple_keyword_rank() -> int:
    """SimpleKeywordStrategy ranks keyword-matching fragments first."""
    strategy = SimpleKeywordStrategy()
    strategy.set_query("async IO")

    frags = [
        _make_frag("project://app/io.py", "Use async IO pattern"),
        _make_frag("project://app/main.py", "Main entry point"),
        _make_frag("project://app/net.py", "Async network handler"),
    ]
    budget = ContextBudget(max_tokens=1000, reserved_tokens=0)

    result = list(strategy.assemble(frags, budget))
    if result[0].uko_node != "project://app/io.py":
        print(f"FAIL: expected io.py first, got {result[0].uko_node}")
        return 1
    print("context-strategies-ok: simple-keyword-rank")
    return 0


def _cmd_simple_keyword_budget() -> int:
    """SimpleKeywordStrategy respects token budget."""
    strategy = SimpleKeywordStrategy()
    strategy.set_query("hello")

    frags = [
        _make_frag("project://app/a.py", "hello world", tokens=100),
        _make_frag("project://app/b.py", "hello there", tokens=100),
        _make_frag("project://app/c.py", "hello again", tokens=100),
    ]
    budget = ContextBudget(max_tokens=250, reserved_tokens=0)

    result = list(strategy.assemble(frags, budget))
    if len(result) != 2:
        print(f"FAIL: expected 2 fragments, got {len(result)}")
        return 1
    print("context-strategies-ok: simple-keyword-budget")
    return 0


def _cmd_semantic_embedding_rank() -> int:
    """SemanticEmbeddingStrategy ranks by word similarity."""
    strategy = SemanticEmbeddingStrategy()
    strategy.set_query("database connection pool")

    frags = [
        _make_frag("project://app/db.py", "Database connection pool manager"),
        _make_frag("project://app/io.py", "File input output handler"),
        _make_frag("project://app/sql.py", "SQL database query executor"),
    ]
    budget = ContextBudget(max_tokens=1000, reserved_tokens=0)

    result = list(strategy.assemble(frags, budget))
    if result[0].uko_node != "project://app/db.py":
        print(f"FAIL: expected db.py first, got {result[0].uko_node}")
        return 1
    print("context-strategies-ok: semantic-embedding-rank")
    return 0


def _cmd_semantic_embedding_filter() -> int:
    """SemanticEmbeddingStrategy filters unrelated fragments."""
    strategy = SemanticEmbeddingStrategy()
    strategy.set_query("quantum computing")

    frags = [
        _make_frag("project://app/db.py", "database handler"),
        _make_frag("project://app/io.py", "file io module"),
    ]
    budget = ContextBudget(max_tokens=1000, reserved_tokens=0)

    result = list(strategy.assemble(frags, budget))
    if len(result) != 0:
        print(f"FAIL: expected 0 fragments, got {len(result)}")
        return 1
    print("context-strategies-ok: semantic-embedding-filter")
    return 0


def _cmd_breadth_depth_rank() -> int:
    """BreadthDepthNavigatorStrategy prioritises near focus."""
    strategy = BreadthDepthNavigatorStrategy()
    strategy.set_focus(["project://app/io.py"])

    frags = [
        _make_frag("project://app/io.py", "io module", depth=5),
        _make_frag("project://app/main.py", "main entry", score=0.9, depth=3),
        _make_frag("project://other/lib.py", "library", score=0.7, depth=9),
    ]
    budget = ContextBudget(max_tokens=1000, reserved_tokens=0)

    result = list(strategy.assemble(frags, budget))
    if result[0].uko_node != "project://app/io.py":
        print(f"FAIL: expected io.py first, got {result[0].uko_node}")
        return 1
    print("context-strategies-ok: breadth-depth-rank")
    return 0


def _cmd_breadth_depth_budget() -> int:
    """BreadthDepthNavigatorStrategy respects budget."""
    strategy = BreadthDepthNavigatorStrategy()
    strategy.set_focus(["project://app"])

    frags = [
        _make_frag("project://app/a.py", "alpha", score=0.9, tokens=100, depth=5),
        _make_frag("project://app/b.py", "beta", score=0.7, tokens=100, depth=5),
        _make_frag("project://app/c.py", "gamma", score=0.5, tokens=100, depth=5),
    ]
    budget = ContextBudget(max_tokens=250, reserved_tokens=0)

    result = list(strategy.assemble(frags, budget))
    if len(result) != 2:
        print(f"FAIL: expected 2 fragments, got {len(result)}")
        return 1
    print("context-strategies-ok: breadth-depth-budget")
    return 0


def _cmd_pipeline_register() -> int:
    """Register all batch 1 strategies with pipeline."""
    pipeline = ACMSPipeline()
    pipeline.register_strategy("simple-keyword", SimpleKeywordStrategy())
    pipeline.register_strategy("semantic-embedding", SemanticEmbeddingStrategy())
    pipeline.register_strategy(
        "breadth-depth-navigator", BreadthDepthNavigatorStrategy()
    )

    registered = pipeline._strategies
    for name in ("simple-keyword", "semantic-embedding", "breadth-depth-navigator"):
        if name not in registered:
            print(f"FAIL: strategy '{name}' not registered")
            return 1
    print("context-strategies-ok: pipeline-register")
    return 0


def _cmd_can_handle() -> int:
    """Verify can_handle confidence values for all strategies."""
    sk = SimpleKeywordStrategy()
    se = SemanticEmbeddingStrategy()
    bd = BreadthDepthNavigatorStrategy()

    # SimpleKeyword always returns 0.3
    if abs(sk.can_handle({"query": "test"}) - 0.3) > 1e-6:
        print("FAIL: SimpleKeyword can_handle != 0.3")
        return 1

    # SemanticEmbedding returns 0.6 with query
    if abs(se.can_handle({"query": "test"}) - 0.6) > 1e-6:
        print("FAIL: SemanticEmbedding can_handle with query != 0.6")
        return 1

    # SemanticEmbedding returns 0.1 without query
    if abs(se.can_handle({}) - 0.1) > 1e-6:
        print("FAIL: SemanticEmbedding can_handle without query != 0.1")
        return 1

    # BreadthDepth returns 0.85 with focus
    if abs(bd.can_handle({"focus": ["project://app"]}) - 0.85) > 1e-6:
        print("FAIL: BreadthDepth can_handle with focus != 0.85")
        return 1

    # BreadthDepth returns 0.2 without focus
    if abs(bd.can_handle({}) - 0.2) > 1e-6:
        print("FAIL: BreadthDepth can_handle without focus != 0.2")
        return 1

    print("context-strategies-ok: can-handle")
    return 0


def _cmd_capabilities() -> int:
    """Verify capability flags for all strategies."""
    sk = SimpleKeywordStrategy()
    se = SemanticEmbeddingStrategy()
    bd = BreadthDepthNavigatorStrategy()

    if sk.capabilities.supports_semantic_search:
        print("FAIL: SimpleKeyword should not support semantic search")
        return 1
    if not se.capabilities.supports_semantic_search:
        print("FAIL: SemanticEmbedding should support semantic search")
        return 1
    if not bd.capabilities.supports_graph_navigation:
        print("FAIL: BreadthDepth should support graph navigation")
        return 1
    if sk.name != "simple-keyword":
        print(f"FAIL: SimpleKeyword name = {sk.name}")
        return 1
    if se.name != "semantic-embedding":
        print(f"FAIL: SemanticEmbedding name = {se.name}")
        return 1
    if bd.name != "breadth-depth-navigator":
        print(f"FAIL: BreadthDepth name = {bd.name}")
        return 1

    print("context-strategies-ok: capabilities")
    return 0


# ---------------------------------------------------------------------------
# Dispatch
# ---------------------------------------------------------------------------

_COMMANDS: dict[str, object] = {
    "simple-keyword-rank": _cmd_simple_keyword_rank,
    "simple-keyword-budget": _cmd_simple_keyword_budget,
    "semantic-embedding-rank": _cmd_semantic_embedding_rank,
    "semantic-embedding-filter": _cmd_semantic_embedding_filter,
    "breadth-depth-rank": _cmd_breadth_depth_rank,
    "breadth-depth-budget": _cmd_breadth_depth_budget,
    "pipeline-register": _cmd_pipeline_register,
    "can-handle": _cmd_can_handle,
    "capabilities": _cmd_capabilities,
}


def main() -> int:
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <command>")
        print(f"Commands: {', '.join(sorted(_COMMANDS))}")
        return 1

    cmd = sys.argv[1]
    handler = _COMMANDS.get(cmd)
    if handler is None:
        print(f"Unknown command: {cmd}")
        return 1

    return handler()  # type: ignore[operator]


if __name__ == "__main__":
    sys.exit(main())