cleveragents-core/robot/helper_context_analysis.py

#!/usr/bin/env python3
"""
Helper script for Context Analysis Agent Robot Framework tests.

This script provides test operations for verifying the ContextAnalysisAgent
workflow functionality.
"""

import sys
import tempfile
from pathlib import Path
from typing import Any

# Add src to path
src_dir = Path(__file__).parent.parent / "src"
sys.path.insert(0, str(src_dir))

from langchain_community.llms import FakeListLLM  # noqa: E402

from cleveragents.agents.context_analysis import (  # noqa: E402
    ContextAnalysisAgent,
    ContextAnalysisState,
)


def test_nodes() -> None:
    """Test that the workflow graph contains all expected nodes."""
    try:
        agent = ContextAnalysisAgent(
            llm=FakeListLLM(
                responses=[
                    "Dependencies: ['os']",
                    "Relevance: High",
                    "Summary: test",
                ]
            )
        )

        # Get the nodes from the graph
        graph = agent.graph
        # nodes is a dict-like attribute of StateGraph
        nodes = getattr(graph, "nodes", {})

        expected_nodes = [
            "load_files",
            "analyze_dependencies",
            "chunk_documents",
            "score_relevance",
            "summarize_context",
        ]

        missing_nodes = [node for node in expected_nodes if node not in nodes]

        if missing_nodes:
            print(f"FAILURE: Missing nodes: {missing_nodes}")
            sys.exit(1)

        print("SUCCESS: All nodes present")
        print(f"Nodes: {list(nodes.keys())}")

    except Exception as e:
        print(f"FAILURE: {e}")
        sys.exit(1)


def test_load_files() -> None:
    """Test file loading functionality."""
    try:
        agent = ContextAnalysisAgent(
            llm=FakeListLLM(
                responses=[
                    "Dependencies: ['os']",
                    "Relevance: High",
                    "Summary: test",
                ]
            )
        )

        # Create a temporary test file
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
            f.write("# Test file\nimport os\nimport sys\n\nprint('hello')\n")
            temp_file = f.name

        try:
            # Create initial state
            state: ContextAnalysisState = {
                "file_paths": [temp_file],
                "documents": [],
                "dependencies": {},
                "summary": "",
                "relevance_scores": {},
                "chunks": [],
                "error": None,
            }

            # Invoke the workflow with required config for checkpointer
            config = {"configurable": {"thread_id": "test_load_files"}}
            result = agent.invoke(state, config)

            # Check that documents were loaded
            if not result.get("documents"):
                print("FAILURE: No documents loaded")
                sys.exit(1)

            print("SUCCESS: Loaded documents")
            print(f"Document count: {len(result['documents'])}")

        finally:
            # Clean up temp file
            Path(temp_file).unlink(missing_ok=True)

    except Exception as e:
        print(f"FAILURE: {e}")
        sys.exit(1)


def test_missing_file() -> None:
    """Test error handling for missing files."""
    try:
        agent = ContextAnalysisAgent(
            llm=FakeListLLM(
                responses=[
                    "Dependencies: ['os']",
                    "Relevance: High",
                    "Summary: test",
                ]
            )
        )

        # Create state with non-existent file
        state: ContextAnalysisState = {
            "file_paths": ["/nonexistent/file.py"],
            "documents": [],
            "dependencies": {},
            "summary": "",
            "relevance_scores": {},
            "chunks": [],
            "error": None,
        }

        # Invoke the workflow with required config for checkpointer
        config = {"configurable": {"thread_id": "test_missing_file"}}
        result = agent.invoke(state, config)

        # Check that error was set
        error_msg = result.get("error")
        if error_msg is None:
            print("FAILURE: No error reported for missing file")
            sys.exit(1)

        if "not found" not in str(error_msg).lower():
            print(f"FAILURE: Unexpected error message: {result['error']}")
            sys.exit(1)

        print("SUCCESS: Error handled")
        print(f"Error: {result['error']}")

    except Exception as e:
        print(f"FAILURE: {e}")
        sys.exit(1)


def test_invoke() -> None:
    """Test complete workflow execution via invoke."""
    try:
        agent = ContextAnalysisAgent(
            llm=FakeListLLM(
                responses=[
                    "Dependencies: ['os']",
                    "Relevance: High",
                    "Summary: test",
                ]
            )
        )

        # Create a temporary test file
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
            f.write("# Test file\nimport os\nimport sys\n\ndef test():\n    pass\n")
            temp_file = f.name

        try:
            # Create initial state
            state: ContextAnalysisState = {
                "file_paths": [temp_file],
                "documents": [],
                "dependencies": {},
                "summary": "",
                "relevance_scores": {},
                "chunks": [],
                "error": None,
            }

            # Invoke the workflow with required config for checkpointer
            config = {"configurable": {"thread_id": "test_invoke"}}
            result = agent.invoke(state, config)

            # Check that all expected fields are populated
            if not result.get("documents"):
                print("FAILURE: No documents in result")
                sys.exit(1)

            if not result.get("dependencies"):
                print("FAILURE: No dependencies in result")
                sys.exit(1)

            if not result.get("chunks"):
                print("FAILURE: No chunks in result")
                sys.exit(1)

            if not result.get("relevance_scores"):
                print("FAILURE: No relevance scores in result")
                sys.exit(1)

            if not result.get("summary"):
                print("FAILURE: No summary in result")
                sys.exit(1)

            print("SUCCESS: Workflow completed")
            print(f"Documents: {len(result['documents'])}")
            print(f"Dependencies: {result['dependencies']}")
            print(f"Chunks: {len(result['chunks'])}")
            print(f"Relevance scores: {result['relevance_scores']}")
            print(f"Summary: {result['summary'][:100]}...")

        finally:
            # Clean up temp file
            Path(temp_file).unlink(missing_ok=True)

    except Exception as e:
        print(f"FAILURE: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)


def test_streaming() -> None:
    """Test streaming workflow execution."""
    try:
        agent = ContextAnalysisAgent(
            llm=FakeListLLM(
                responses=[
                    "Dependencies: ['os']",
                    "Relevance: High",
                    "Summary: test",
                ]
            )
        )

        # Create a temporary test file
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
            f.write("# Test file\nimport os\n\nprint('test')\n")
            temp_file = f.name

        try:
            # Create initial state
            state: ContextAnalysisState = {
                "file_paths": [temp_file],
                "documents": [],
                "dependencies": {},
                "summary": "",
                "relevance_scores": {},
                "chunks": [],
                "error": None,
            }

            # Stream the workflow with required config for checkpointer
            config = {"configurable": {"thread_id": "test_streaming"}}
            updates: list[dict[str, Any]] = []
            for event in agent.stream(state, config):
                updates.append(event)

            # Check that we got stream updates
            if len(updates) == 0:
                print("FAILURE: No stream updates received")
                sys.exit(1)

            print("SUCCESS: Received stream updates")
            print(f"Update count: {len(updates)}")

        finally:
            # Clean up temp file
            Path(temp_file).unlink(missing_ok=True)

    except Exception as e:
        print(f"FAILURE: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)


def main() -> None:
    """Main entry point for the helper script."""
    if len(sys.argv) < 2:
        print("Usage: helper_context_analysis.py <test_name>")
        print("Available tests: nodes, load_files, missing_file, invoke, streaming")
        sys.exit(1)

    test_name = sys.argv[1]

    test_functions = {
        "nodes": test_nodes,
        "load_files": test_load_files,
        "missing_file": test_missing_file,
        "invoke": test_invoke,
        "streaming": test_streaming,
    }

    test_func = test_functions.get(test_name)
    if test_func is None:
        print(f"Unknown test: {test_name}")
        print(f"Available tests: {', '.join(test_functions.keys())}")
        sys.exit(1)

    test_func()


if __name__ == "__main__":
    main()