cleveragents-core/features/mocks/langchain_mock_provider.py

"""LangChain-based mock AI provider for testing purposes only.

This mock implementation uses LangChain's FakeListLLM for deterministic testing,
following the mock placement rule that all mocks must exist in the features/
directory, never in production code (ADR-022).
"""

import os
from collections.abc import Callable, Iterator
from typing import Any

from langchain_community.llms import FakeListLLM
from langchain_core.language_models import BaseLLM
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate

from cleveragents.domain.models.core import (
    Change,
    Context,
    OperationType,
    Plan,
    Project,
)
from cleveragents.domain.providers.ai_provider import (
    ActorInvocationContext,
    ProviderResponse,
)


class FailingLLM(BaseLLM):
    """A LangChain LLM that always fails for testing error handling."""

    @property
    def _llm_type(self) -> str:
        return "failing"

    def _call(
        self,
        prompt: str,
        stop: list[str] | None = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> str:
        """Always raise an exception."""
        raise Exception("Mock AI provider configured to fail")

    async def _acall(
        self,
        prompt: str,
        stop: list[str] | None = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> str:
        """Always raise an exception."""
        raise Exception("Mock AI provider configured to fail")


def get_mock_llm() -> BaseLLM:
    """Get the appropriate mock LLM based on environment variables.

    Returns:
        FailingLLM if CLEVERAGENTS_MOCK_SHOULD_FAIL is set,
        otherwise a standard FakeListLLM for testing.
    """
    if os.environ.get("CLEVERAGENTS_MOCK_SHOULD_FAIL") == "true":
        return FailingLLM()

    # Return standard mock with deterministic responses
    return FakeListLLM(
        responses=[
            "Requirements: Add error handling with try-except blocks",
            "Generated code with proper error handling implementation",
            "Validation passed: Code follows best practices",
        ]
    )


class LangChainMockProvider:
    """LangChain-based mock AI provider that generates deterministic test responses.

    This provider uses LangChain's FakeListLLM for deterministic testing,
    avoiding calls to real AI APIs during tests.
    """

    def __init__(self, model_id: str = "langchain-mock-gpt-4"):
        """Initialize the LangChain mock provider.

        Args:
            model_id: Mock model identifier
        """
        self._model_id = model_id
        self._name = "LangChainMockProvider"

        # Define deterministic responses for different prompt patterns
        self._responses = [
            """# Error handling code
def handle_error(error):
    \"\"\"Handle errors gracefully.\"\"\"
    import logging
    logging.error(f"Error occurred: {error}")
    return None
""",
            """# Test file
import pytest

def test_example():
    \"\"\"Test example functionality.\"\"\"
    assert True

def test_error_handling():
    \"\"\"Test error handling.\"\"\"
    with pytest.raises(ValueError):
        raise ValueError("Test error")
""",
            """# Refactored code
class RefactoredClass:
    \"\"\"Improved implementation.\"\"\"

    def __init__(self):
        self.state = "initialized"

    def process(self, data):
        \"\"\"Process data with better error handling.\"\"\"
        if not data:
            return None
        return f"Processed: {data}"
""",
            """# Generated code
def main():
    \"\"\"Main entry point.\"\"\"
    print("Hello from CleverAgents!")
    return 0

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

        # Initialize FakeListLLM with responses
        self._llm = FakeListLLM(responses=self._responses)

        # Create a simple prompt template
        self._prompt_template = PromptTemplate(
            input_variables=["instruction", "context"],
            template="""
You are a code generation assistant. Generate code based on the following:

Instruction: {instruction}

Context Files:
{context}

Generated Code:
""",
        )

        # Create chain
        self._chain = self._prompt_template | self._llm | StrOutputParser()

    def generate_changes(
        self,
        project: Project,
        plan: Plan,
        contexts: list[Context],
        actor_context: ActorInvocationContext | None = None,
        progress_callback: Callable[[int], None] | None = None,
    ) -> ProviderResponse:
        """Generate mock code changes based on the plan using LangChain.

        Args:
            project: The project being modified
            plan: The plan containing instructions
            contexts: List of context files to consider
            progress_callback: Optional callback for progress updates (0-100)

        Returns:
            ProviderResponse containing mock changes
        """
        # Check if we should fail for testing
        import os

        if os.environ.get("CLEVERAGENTS_MOCK_SHOULD_FAIL") == "true":
            raise Exception("Mock AI provider configured to fail")

        # Check if we should generate invalid code for testing
        if os.environ.get("CLEVERAGENTS_MOCK_INVALID_CODE") == "true":
            # Generate syntactically invalid Python code
            generated_code = (
                "def invalid_function(\n    # Missing closing paren and body"
            )
            prompt_lower = ""
        else:
            # Simulate progress
            if progress_callback:
                for i in range(0, 101, 20):
                    progress_callback(i)

            # Prepare context for the prompt
            context_str = ""
            if contexts:
                context_str = "\n".join(
                    f"File: {ctx.path}\nContent: {ctx.content[:200]}..."
                    for ctx in contexts[:3]  # Limit to first 3 for brevity
                )

            # Use LangChain to generate response
            prompt_lower = plan.prompt.lower() if plan.prompt else ""

            try:
                # Invoke the chain with the prompt
                generated_code = self._chain.invoke(
                    {
                        "instruction": plan.prompt or "Generate example code",
                        "context": context_str or "No context files provided",
                    }
                )
            except Exception:
                # Fallback to default if chain fails
                generated_code = (
                    "# Default generated code\nprint('Hello from CleverAgents!')\n"
                )

        # Determine file path and operation based on prompt
        if "error handling" in prompt_lower or "exception" in prompt_lower:
            file_path = "error_handler.py"
            operation = OperationType.CREATE
        elif "test" in prompt_lower:
            file_path = "test_example.py"
            operation = OperationType.CREATE
        elif "refactor" in prompt_lower or "improve" in prompt_lower:
            if contexts and len(contexts) > 0:
                file_path = contexts[0].path
                operation = OperationType.MODIFY
            else:
                file_path = "refactored.py"
                operation = OperationType.CREATE
        else:
            file_path = "generated.py"
            operation = OperationType.CREATE

        # Create change object
        plan_id = plan.id if plan.id else 0  # Use 0 as placeholder for tests

        changes = [
            Change(
                id=None,
                plan_id=plan_id,
                file_path=file_path,
                operation=operation,
                original_content=contexts[0].content
                if contexts and operation == OperationType.MODIFY
                else None,
                new_content=generated_code,
                applied=False,
                applied_at=None,
                new_path=None,
            )
        ]

        # Calculate approximate token count (rough estimate)
        token_count = len(generated_code.split()) * 1.3  # Rough approximation

        return ProviderResponse(
            changes=changes,
            model_used=self._model_id,
            token_count=int(token_count),
            error_message=None,
        )

    def stream_changes(
        self,
        project: Project,
        plan: Plan,
        contexts: list[Context],
        actor_context: ActorInvocationContext | None = None,
        progress_callback: Callable[[int], None] | None = None,
    ) -> Iterator[dict[str, object]]:
        """Stream mock workflow events for CLI coverage."""

        response = self.generate_changes(
            project,
            plan,
            contexts,
            progress_callback=progress_callback,
        )

        for node in (
            "load_context",
            "analyze_requirements",
            "generate_plan",
            "validate",
        ):
            payload: dict[str, object] = {}
            if node == "generate_plan":
                payload = {
                    "status": "completed",
                    "change_count": len(response.changes),
                }
            elif node == "validate":
                payload = {"status": "completed"}
            yield {node: payload}

        yield {"__end__": {"response": response}}

    @property
    def name(self) -> str:
        """Get the provider name."""
        return self._name

    @property
    def model_id(self) -> str:
        """Get the model identifier."""
        return self._model_id