TemporaryFailureModel Class — langchain Architecture

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_model_retry.py lines 26–60

class TemporaryFailureModel(FakeToolCallingModel):
    """Model that fails a certain number of times before succeeding."""

    fail_count: int = Field(default=0)
    attempt: int = Field(default=0)

    def _generate(
        self,
        messages: list[BaseMessage],
        stop: list[str] | None = None,
        run_manager: CallbackManagerForLLMRun | None = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Execute the model.

        Args:
            messages: Input messages.
            stop: Optional stop sequences.
            run_manager: Optional callback manager.
            **kwargs: Additional keyword arguments.

        Returns:
            ChatResult with success message if attempt >= fail_count.

        Raises:
            ValueError: If attempt < fail_count.
        """
        self.attempt += 1
        if self.attempt <= self.fail_count:
            msg = f"Temporary failure {self.attempt}"
            raise ValueError(msg)
        # Return success message
        ai_msg = AIMessage(content=f"Success after {self.attempt} attempts", id=str(self.index))
        self.index += 1
        return ChatResult(generations=[ChatGeneration(message=ai_msg)])